feat: 重组项目结构并完善运行管理
This commit is contained in:
@@ -0,0 +1,49 @@
|
||||
# Changelog
|
||||
|
||||
本项目遵循语义化版本。0.2.1 是包含公开 API 变更的 breaking release。
|
||||
|
||||
## Unreleased
|
||||
|
||||
### Added
|
||||
|
||||
- BEP-51 `sample_infohashes` 主动发现、按节点 interval/退避的采样 actor,以及有界 Hash
|
||||
去重。
|
||||
- 采样 Hash 优先查询来源节点,再通过迭代式 `get_peers` 补充 Peer。
|
||||
|
||||
### Changed
|
||||
|
||||
- Metadata 最大 info 字典大小改为可通过 `max_metadata_size_bytes` 配置,并拒绝负数文件大小和总大小溢出。
|
||||
- 主动 `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% 保底预算。
|
||||
@@ -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"
|
||||
@@ -0,0 +1,272 @@
|
||||
# dht-crawler
|
||||
|
||||
[](https://crates.io/crates/dht-crawler)
|
||||
[](https://docs.rs/dht-crawler)
|
||||
[](../../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,
|
||||
max_metadata_size_bytes: 10 * 1024 * 1024,
|
||||
..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 查询
|
||||
- `new_node_sample_percent` 大于零时按节点地址稳定分流并优先直接采样 避免同一次发现同时执行 `find_node`
|
||||
- 新节点采样通道满时回退到抓取池 不会依靠无限队列维持吞吐
|
||||
- 带首选节点的 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)
|
||||
@@ -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 延迟。
|
||||
@@ -0,0 +1,151 @@
|
||||
// 负责演示基础 DHT 爬虫的配置启动监控和安全停止方式
|
||||
|
||||
#[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::<Vec<_>>()
|
||||
// .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(())
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
// 负责校验解析和规范化 DHT 节点网络地址
|
||||
|
||||
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()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
// 负责管理 DHT 引导节点解析重试退避和准入调度
|
||||
|
||||
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<Instant>,
|
||||
}
|
||||
|
||||
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<Instant>,
|
||||
last_success: Option<Instant>,
|
||||
fail_count: u32,
|
||||
backoff_until: Option<Instant>,
|
||||
}
|
||||
|
||||
pub(crate) struct BootstrapSourcePool {
|
||||
pub(crate) hosts: Vec<String>,
|
||||
states: HashMap<SocketAddr, BootstrapSourceState>,
|
||||
backoff_base: Duration,
|
||||
backoff_max: Duration,
|
||||
}
|
||||
|
||||
impl BootstrapSourcePool {
|
||||
pub(crate) fn new(hosts: Vec<String>, backoff_base: Duration, backoff_max: Duration) -> Self {
|
||||
Self {
|
||||
hosts,
|
||||
states: HashMap::new(),
|
||||
backoff_base,
|
||||
backoff_max,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn select(
|
||||
&mut self,
|
||||
candidates: Vec<SocketAddr>,
|
||||
max_nodes: usize,
|
||||
now: Instant,
|
||||
) -> Vec<SocketAddr> {
|
||||
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<SocketAddr> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
// 负责实现网络查询共享速率预算和突发容量控制
|
||||
|
||||
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<Mutex<RateBucket>>,
|
||||
}
|
||||
|
||||
impl SharedRateBudget {
|
||||
pub(crate) fn per_second(rate_per_sec: u32, burst: u32, initially_full: bool) -> Self {
|
||||
Self {
|
||||
bucket: Arc::new(Mutex::new(RateBucket::per_second(
|
||||
rate_per_sec,
|
||||
burst,
|
||||
initially_full,
|
||||
Instant::now(),
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn try_take_one(&self, now: Instant) -> bool {
|
||||
self.bucket
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.try_take_one(now)
|
||||
}
|
||||
|
||||
pub(crate) fn refund_one(&self) {
|
||||
self.bucket
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.refund_one();
|
||||
}
|
||||
}
|
||||
|
||||
impl RateBucket {
|
||||
pub(crate) fn per_second(
|
||||
rate_per_sec: u32,
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
// 负责解析抓取配置并生成运行时资源限制
|
||||
|
||||
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<String>,
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,974 @@
|
||||
// 负责协调节点探测 Peer 查找采样发现和 Metadata 调度
|
||||
|
||||
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::sample_infohashes::SampleCandidateRouter;
|
||||
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<SocketAddr>),
|
||||
}
|
||||
|
||||
struct OutboundRequest {
|
||||
key: PendingKey,
|
||||
node: NodeTuple,
|
||||
target: [u8; 20],
|
||||
sender_id: [u8; 20],
|
||||
}
|
||||
|
||||
pub(crate) struct CrawlEngine {
|
||||
config: ResolvedCrawlConfig,
|
||||
priority_tx: mpsc::Sender<PriorityEvent>,
|
||||
discovery_tx: mpsc::Sender<NodeTuple>,
|
||||
/// One-shot handoff used only by `DHTServer::start`; never touched by the crawl hot path.
|
||||
receivers: Mutex<Option<(mpsc::Receiver<PriorityEvent>, mpsc::Receiver<NodeTuple>)>>,
|
||||
pub(crate) snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
pub(crate) node_count: Arc<AtomicUsize>,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
outbound_query_budget: SharedRateBudget,
|
||||
sample_candidate_router: Option<SampleCandidateRouter>,
|
||||
}
|
||||
|
||||
impl CrawlEngine {
|
||||
pub(crate) fn new(
|
||||
config: ResolvedCrawlConfig,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
outbound_query_budget: SharedRateBudget,
|
||||
sample_candidate_router: Option<SampleCandidateRouter>,
|
||||
) -> 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,
|
||||
sample_candidate_router,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn route_discovered(&self, node: NodeTuple) {
|
||||
if self
|
||||
.sample_candidate_router
|
||||
.as_ref()
|
||||
.is_some_and(|router| router.route(node))
|
||||
{
|
||||
return;
|
||||
}
|
||||
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<SocketAddr, Arc<UdpSocket>>,
|
||||
metadata_queue_len: Arc<AtomicUsize>,
|
||||
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<UdpSocket>,
|
||||
mut rx: mpsc::Receiver<OutboundRequest>,
|
||||
priority_tx: mpsc::Sender<PriorityEvent>,
|
||||
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<PriorityEvent>,
|
||||
discovery_rx: mpsc::Receiver<NodeTuple>,
|
||||
priority_tx: mpsc::Sender<PriorityEvent>,
|
||||
egress_v4: Option<mpsc::Sender<OutboundRequest>>,
|
||||
egress_v6: Option<mpsc::Sender<OutboundRequest>>,
|
||||
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
node_count: Arc<AtomicUsize>,
|
||||
metadata_queue_len: Arc<AtomicUsize>,
|
||||
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<PriorityEvent>,
|
||||
discovery_rx: mpsc::Receiver<NodeTuple>,
|
||||
priority_tx: mpsc::Sender<PriorityEvent>,
|
||||
egress_v4: Option<mpsc::Sender<OutboundRequest>>,
|
||||
egress_v6: Option<mpsc::Sender<OutboundRequest>>,
|
||||
pool: NodePool,
|
||||
responsive: ResponsiveReservoir,
|
||||
pending: AHashMap<PendingKey, PendingRequest>,
|
||||
pending_expiry: VecDeque<(Instant, PendingKey)>,
|
||||
subnet_in_flight: AHashMap<SubnetKey, usize>,
|
||||
query_budget: RateBucket,
|
||||
destination_budget: RateBucket,
|
||||
bootstrap_gate: BootstrapGate,
|
||||
bootstrap_pool: BootstrapSourcePool,
|
||||
bootstrap_queue: VecDeque<SocketAddr>,
|
||||
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
node_count: Arc<AtomicUsize>,
|
||||
metadata_queue_len: Arc<AtomicUsize>,
|
||||
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<OutboundRequest>, 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),
|
||||
None,
|
||||
);
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// 负责定义 DHT 基础库的统一错误类型和结果别名
|
||||
|
||||
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<T> = std::result::Result<T, DHTError>;
|
||||
@@ -0,0 +1,310 @@
|
||||
// 负责编码 DHT KRPC 查询响应和紧凑节点数据
|
||||
|
||||
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<UdpSocket>,
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
// 负责导出可复用 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 persistence 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 response_limiter;
|
||||
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;
|
||||
pub use types::{
|
||||
BootstrapOptions, CrawlOptions, DHTOptions, DiscoverySource, FileInfo, HashDiscovered,
|
||||
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, DiscoverySource, FileInfo, HashDiscovered,
|
||||
MetadataFetchCompletion, MetadataFetchCompletionStatus, MetadataOptions, NetMode,
|
||||
NodeTuple, PeerLookupOptions, PeerLookupResult, PoolOptions, RateLimitOptions,
|
||||
SampleInfohashesOptions, SchedulerOptions, TargetOptions, TorrentInfo,
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,740 @@
|
||||
// 负责执行 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<FileInfo>, 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<Instant>,
|
||||
}
|
||||
|
||||
impl ConnectRateLimiter {
|
||||
fn per_second(rate: u32) -> Self {
|
||||
Self {
|
||||
interval: Duration::from_secs_f64(1.0 / f64::from(rate.max(1))),
|
||||
next_start: Mutex::new(Instant::now()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn acquire(&self) {
|
||||
let delay = {
|
||||
let mut next_start = self
|
||||
.next_start
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let now = Instant::now();
|
||||
let reserved = (*next_start).max(now);
|
||||
*next_start = reserved + self.interval;
|
||||
reserved.saturating_duration_since(now)
|
||||
};
|
||||
if !delay.is_zero() {
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
impl PeerFailureReason {
|
||||
fn as_str(self) -> &'static str {
|
||||
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<SocketAddr, PeerFailureEntry>,
|
||||
expiry: VecDeque<(Instant, SocketAddr)>,
|
||||
}
|
||||
|
||||
struct PeerFailureCache {
|
||||
inner: Mutex<PeerFailureCacheInner>,
|
||||
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<PeerFailureReason>, 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,
|
||||
max_metadata_size_bytes: usize,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
peer_failure_cache: Arc<PeerFailureCache>,
|
||||
connect_rate_limiter: Arc<ConnectRateLimiter>,
|
||||
}
|
||||
|
||||
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,
|
||||
10 * 1024 * 1024,
|
||||
200_000,
|
||||
60,
|
||||
DhtRuntimeStats::default(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn new_with_runtime_stats(
|
||||
timeout_secs: u64,
|
||||
max_connects_per_second: u32,
|
||||
max_metadata_size_bytes: usize,
|
||||
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 }),
|
||||
max_metadata_size_bytes: max_metadata_size_bytes.max(1),
|
||||
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<F>(
|
||||
&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<FetchedMetadata, MetadataFetchFailure> {
|
||||
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<u32, Bytes> = 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 as usize > self.max_metadata_size_bytes {
|
||||
#[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<FetchedMetadata> {
|
||||
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())
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.unwrap_or(0);
|
||||
|
||||
let mut total_size = 0_u64;
|
||||
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 = u64::try_from(length).ok()?;
|
||||
total_size = total_size.checked_add(length)?;
|
||||
let parts = file_dict
|
||||
.get(&b"path"[..])
|
||||
.and_then(|value| value.as_list())?;
|
||||
let path = parts
|
||||
.iter()
|
||||
.map(|part| part.as_str())
|
||||
.collect::<Option<Vec<_>>>()?
|
||||
.join("/");
|
||||
file_list.push(FileInfo { path, size: length });
|
||||
}
|
||||
} else if let Some(length) = dict
|
||||
.get(&b"length"[..])
|
||||
.and_then(|value| value.as_integer())
|
||||
{
|
||||
total_size = u64::try_from(length).ok()?;
|
||||
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::*;
|
||||
|
||||
#[test]
|
||||
fn parser_rejects_negative_file_size() {
|
||||
assert!(parse_metadata(b"d6:lengthi-1e4:name1:ae").is_none());
|
||||
}
|
||||
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 * 1024 * 1024, 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 * 1024 * 1024,
|
||||
10,
|
||||
60,
|
||||
DhtRuntimeStats::default(),
|
||||
);
|
||||
assert!(fetcher.verify_handshake(&[7; 20], addr).await);
|
||||
peer.await.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
// 负责生成 DHT 节点标识事务标识和距离目标
|
||||
|
||||
use rand::RngExt;
|
||||
|
||||
pub(crate) type TransactionId = [u8; 8];
|
||||
|
||||
pub(crate) fn transaction_id_from_bytes(bytes: &[u8]) -> Option<TransactionId> {
|
||||
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<u8> {
|
||||
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]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,392 @@
|
||||
// 负责维护有界 DHT 节点池去重淘汰和重访状态
|
||||
|
||||
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<QueuedNode>,
|
||||
queued: AHashSet<SocketAddr>,
|
||||
recent: AHashMap<SocketAddr, Instant>,
|
||||
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<NodeTuple> {
|
||||
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<NodeTuple> {
|
||||
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<QueuedNode> {
|
||||
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<Option<ResponsiveEntry>>,
|
||||
index: AHashMap<SocketAddr, usize>,
|
||||
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<NodeTuple> {
|
||||
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<NodeTuple> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,911 @@
|
||||
// 负责执行有界主动 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::types::{
|
||||
DiscoverySource, HashDiscovered, 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<NodeTuple>,
|
||||
preferred: Option<NodeTuple>,
|
||||
seen_nodes: AHashSet<SocketAddr>,
|
||||
peers: AHashSet<SocketAddr>,
|
||||
queried: usize,
|
||||
outstanding: usize,
|
||||
deadline: Instant,
|
||||
preferred_phase: bool,
|
||||
allow_iterative_fallback: bool,
|
||||
source: DiscoverySource,
|
||||
completion: Option<oneshot::Sender<PeerLookupResult>>,
|
||||
}
|
||||
|
||||
impl LookupState {
|
||||
fn pop_closest(&mut self) -> Option<NodeTuple> {
|
||||
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<NodeTuple>,
|
||||
pub(crate) allow_iterative_fallback: bool,
|
||||
pub(crate) source: DiscoverySource,
|
||||
pub(crate) completion: Option<oneshot::Sender<PeerLookupResult>>,
|
||||
}
|
||||
|
||||
impl PeerLookupRequest {
|
||||
pub(crate) fn new(info_hash: [u8; 20]) -> Self {
|
||||
Self {
|
||||
info_hash,
|
||||
preferred_node: None,
|
||||
allow_iterative_fallback: true,
|
||||
source: DiscoverySource::ActiveLookup,
|
||||
completion: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct PeerLookupHandle {
|
||||
request_tx: mpsc::Sender<PeerLookupRequest>,
|
||||
response_tx: mpsc::Sender<LookupResponse>,
|
||||
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<PeerLookupRequest> {
|
||||
self.request_tx.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn lookup(
|
||||
&self,
|
||||
info_hash: [u8; 20],
|
||||
) -> Result<PeerLookupResult, &'static str> {
|
||||
let (completion, receiver) = oneshot::channel();
|
||||
self.request_tx
|
||||
.send(PeerLookupRequest {
|
||||
info_hash,
|
||||
preferred_node: None,
|
||||
allow_iterative_fallback: true,
|
||||
source: DiscoverySource::ActiveLookup,
|
||||
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<SocketAddr, Arc<UdpSocket>>,
|
||||
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
hash_tx: mpsc::Sender<HashDiscovered>,
|
||||
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<Arc<UdpSocket>>,
|
||||
socket_v6: Option<Arc<UdpSocket>>,
|
||||
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
hash_tx: mpsc::Sender<HashDiscovered>,
|
||||
request_rx: mpsc::Receiver<PeerLookupRequest>,
|
||||
response_rx: mpsc::Receiver<LookupResponse>,
|
||||
request_budget: RateBucket,
|
||||
max_active_lookups: usize,
|
||||
enabled: bool,
|
||||
queued: VecDeque<PeerLookupRequest>,
|
||||
queued_hashes: AHashSet<[u8; 20]>,
|
||||
active_hashes: AHashSet<[u8; 20]>,
|
||||
active: AHashMap<u64, LookupState>,
|
||||
pending: AHashMap<PendingKey, PendingQuery>,
|
||||
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(),
|
||||
allow_iterative_fallback: request.allow_iterative_fallback,
|
||||
source: request.source,
|
||||
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) {
|
||||
if !state.allow_iterative_fallback {
|
||||
self.finish_lookup(lookup_id);
|
||||
return;
|
||||
}
|
||||
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 preferred_without_fallback = pending.preferred_phase && !state.allow_iterative_fallback;
|
||||
let source = state.source;
|
||||
let _ = state;
|
||||
|
||||
for peer in discovered {
|
||||
let event = HashDiscovered {
|
||||
info_hash: hash.clone(),
|
||||
peer_addr: peer,
|
||||
source,
|
||||
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 preferred_without_fallback {
|
||||
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();
|
||||
let mut finish_without_fallback = 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;
|
||||
if state.allow_iterative_fallback {
|
||||
state.deadline = now + LOOKUP_TIMEOUT;
|
||||
self.runtime_stats.peer_lookup_fallback();
|
||||
} else {
|
||||
finish_without_fallback.insert(pending.lookup_id);
|
||||
}
|
||||
}
|
||||
affected.insert(pending.lookup_id);
|
||||
}
|
||||
self.runtime_stats.peer_lookup_timeout();
|
||||
}
|
||||
for lookup_id in finish_without_fallback {
|
||||
affected.remove(&lookup_id);
|
||||
self.finish_lookup(lookup_id);
|
||||
}
|
||||
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,
|
||||
}),
|
||||
allow_iterative_fallback: true,
|
||||
source: DiscoverySource::SampleDirect,
|
||||
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());
|
||||
assert_eq!(event.source, DiscoverySource::SampleDirect);
|
||||
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,
|
||||
}),
|
||||
allow_iterative_fallback: true,
|
||||
source: DiscoverySource::SampleDirect,
|
||||
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();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preferred_node_without_peers_can_finish_without_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(),
|
||||
},
|
||||
);
|
||||
let (completion, result_rx) = oneshot::channel();
|
||||
handle
|
||||
.request_tx
|
||||
.send(PeerLookupRequest {
|
||||
info_hash: [3; 20],
|
||||
preferred_node: Some(NodeTuple {
|
||||
id: [9; 20],
|
||||
addr: preferred_addr,
|
||||
}),
|
||||
allow_iterative_fallback: false,
|
||||
source: DiscoverySource::SampleDirect,
|
||||
completion: Some(completion),
|
||||
})
|
||||
.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,
|
||||
},
|
||||
);
|
||||
|
||||
let result = tokio::time::timeout(Duration::from_secs(1), result_rx)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(result.peers.is_empty());
|
||||
assert_eq!(result.queries, 1);
|
||||
assert!(
|
||||
tokio::time::timeout(
|
||||
Duration::from_millis(200),
|
||||
fallback_socket.recv_from(&mut buffer),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let snapshot = stats.snapshot();
|
||||
assert_eq!(snapshot.peer_lookup_queries, 1);
|
||||
assert_eq!(snapshot.peer_lookup_fallbacks, 0);
|
||||
shutdown.cancel();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
// 负责定义 DHT KRPC 消息查询参数和响应字段的反序列化模型
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
/// Decoded KRPC envelope.
|
||||
pub struct DhtMessage {
|
||||
/// Transaction ID bytes.
|
||||
pub t: serde_bytes::ByteBuf,
|
||||
/// Message kind (`q`, `r`, or `e`).
|
||||
pub y: String,
|
||||
/// Query method when `y == q`.
|
||||
pub q: Option<String>,
|
||||
/// Query arguments.
|
||||
pub a: Option<DhtArgs>,
|
||||
/// Response dictionary.
|
||||
pub r: Option<DhtResponse>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
/// Supported BEP-5 query arguments.
|
||||
pub struct DhtArgs {
|
||||
/// Sender node ID.
|
||||
pub id: Option<serde_bytes::ByteBuf>,
|
||||
/// find_node target ID.
|
||||
pub target: Option<serde_bytes::ByteBuf>,
|
||||
/// get_peers/announce InfoHash.
|
||||
pub info_hash: Option<serde_bytes::ByteBuf>,
|
||||
/// announce validation token.
|
||||
pub token: Option<serde_bytes::ByteBuf>,
|
||||
/// Explicit announced Peer port.
|
||||
pub port: Option<u16>,
|
||||
/// Non-zero means use the UDP source port.
|
||||
pub implied_port: Option<u8>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
/// Supported BEP-5 response fields.
|
||||
pub struct DhtResponse {
|
||||
#[serde(default)]
|
||||
/// Responder node ID.
|
||||
pub id: Option<serde_bytes::ByteBuf>,
|
||||
#[serde(default)]
|
||||
/// Compact IPv4 node tuples.
|
||||
pub nodes: Option<serde_bytes::ByteBuf>,
|
||||
#[serde(default)]
|
||||
/// Compact IPv6 node tuples.
|
||||
pub nodes6: Option<serde_bytes::ByteBuf>,
|
||||
#[serde(default)]
|
||||
/// Compact Peer endpoints returned by `get_peers`.
|
||||
pub values: Option<Vec<serde_bytes::ByteBuf>>,
|
||||
#[serde(default)]
|
||||
/// BEP-51 concatenated 20-byte sampled InfoHashes.
|
||||
pub samples: Option<serde_bytes::ByteBuf>,
|
||||
#[serde(default)]
|
||||
/// BEP-51 estimated number of InfoHashes held by the responder.
|
||||
pub num: Option<u64>,
|
||||
#[serde(default)]
|
||||
/// BEP-51 requested delay before sampling this node again, in seconds.
|
||||
pub interval: Option<u64>,
|
||||
}
|
||||
@@ -0,0 +1,268 @@
|
||||
// 负责限制 DHT 响应的包速率字节速率单来源速率和优先保留预算
|
||||
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
net::SocketAddr,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use ahash::AHashMap;
|
||||
|
||||
use crate::budget::RateBucket;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct SourceResponseWindow {
|
||||
started_at: Instant,
|
||||
last_seen: Instant,
|
||||
count: u32,
|
||||
}
|
||||
|
||||
pub(crate) struct WorkerResponseLimiter {
|
||||
regular_packets: RateBucket,
|
||||
regular_bytes: RateBucket,
|
||||
priority_packets: RateBucket,
|
||||
priority_bytes: RateBucket,
|
||||
per_source_rate: u32,
|
||||
sources: AHashMap<SocketAddr, SourceResponseWindow>,
|
||||
source_expiry: VecDeque<(Instant, SocketAddr)>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum ResponsePermit {
|
||||
Regular,
|
||||
PriorityReserve,
|
||||
Rejected,
|
||||
}
|
||||
|
||||
impl WorkerResponseLimiter {
|
||||
pub(crate) fn new(
|
||||
packet_rate: u32,
|
||||
byte_rate: u64,
|
||||
per_source_rate: u32,
|
||||
now: Instant,
|
||||
) -> Self {
|
||||
let byte_rate = byte_rate.min(u32::MAX as u64) as u32;
|
||||
let priority_packet_rate = reserve_quota(packet_rate);
|
||||
let priority_byte_rate = reserve_quota(byte_rate);
|
||||
let packet_rate = packet_rate.saturating_sub(priority_packet_rate);
|
||||
let byte_rate = byte_rate.saturating_sub(priority_byte_rate);
|
||||
Self {
|
||||
regular_packets: RateBucket::per_second(
|
||||
packet_rate,
|
||||
packet_rate.div_ceil(5).max(1),
|
||||
true,
|
||||
now,
|
||||
),
|
||||
regular_bytes: RateBucket::per_second(
|
||||
byte_rate,
|
||||
byte_rate.div_ceil(5).max(512),
|
||||
true,
|
||||
now,
|
||||
),
|
||||
priority_packets: RateBucket::per_second(
|
||||
priority_packet_rate,
|
||||
priority_packet_rate.div_ceil(5).max(1),
|
||||
true,
|
||||
now,
|
||||
),
|
||||
priority_bytes: RateBucket::per_second(
|
||||
priority_byte_rate,
|
||||
priority_byte_rate.div_ceil(5).max(512),
|
||||
true,
|
||||
now,
|
||||
),
|
||||
per_source_rate,
|
||||
sources: AHashMap::new(),
|
||||
source_expiry: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn acquire(
|
||||
&mut self,
|
||||
addr: SocketAddr,
|
||||
encoded_len: usize,
|
||||
is_priority: bool,
|
||||
now: Instant,
|
||||
) -> ResponsePermit {
|
||||
self.expire_sources(now);
|
||||
if Self::take_budget(
|
||||
&mut self.regular_packets,
|
||||
&mut self.regular_bytes,
|
||||
encoded_len,
|
||||
now,
|
||||
) {
|
||||
if self.acquire_source_slot(addr, now) {
|
||||
return ResponsePermit::Regular;
|
||||
}
|
||||
Self::refund_budget(
|
||||
&mut self.regular_packets,
|
||||
&mut self.regular_bytes,
|
||||
encoded_len,
|
||||
);
|
||||
return ResponsePermit::Rejected;
|
||||
}
|
||||
if is_priority
|
||||
&& Self::take_budget(
|
||||
&mut self.priority_packets,
|
||||
&mut self.priority_bytes,
|
||||
encoded_len,
|
||||
now,
|
||||
)
|
||||
{
|
||||
if self.acquire_source_slot(addr, now) {
|
||||
return ResponsePermit::PriorityReserve;
|
||||
}
|
||||
Self::refund_budget(
|
||||
&mut self.priority_packets,
|
||||
&mut self.priority_bytes,
|
||||
encoded_len,
|
||||
);
|
||||
}
|
||||
ResponsePermit::Rejected
|
||||
}
|
||||
|
||||
fn take_budget(
|
||||
packets: &mut RateBucket,
|
||||
bytes: &mut RateBucket,
|
||||
encoded_len: usize,
|
||||
now: Instant,
|
||||
) -> bool {
|
||||
if !packets.try_take_one(now) {
|
||||
return false;
|
||||
}
|
||||
if !bytes.try_take_exact(encoded_len, now) {
|
||||
packets.refund_one();
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn refund_budget(packets: &mut RateBucket, bytes: &mut RateBucket, encoded_len: usize) {
|
||||
packets.refund_one();
|
||||
bytes.refund(encoded_len);
|
||||
}
|
||||
|
||||
fn acquire_source_slot(&mut self, addr: SocketAddr, now: Instant) -> bool {
|
||||
if self.per_source_rate == 0 {
|
||||
return false;
|
||||
}
|
||||
let entry = self.sources.entry(addr).or_insert(SourceResponseWindow {
|
||||
started_at: now,
|
||||
last_seen: now,
|
||||
count: 0,
|
||||
});
|
||||
if now
|
||||
.checked_duration_since(entry.started_at)
|
||||
.unwrap_or_default()
|
||||
>= Duration::from_secs(1)
|
||||
{
|
||||
entry.started_at = now;
|
||||
entry.count = 0;
|
||||
}
|
||||
if entry.count >= self.per_source_rate {
|
||||
return false;
|
||||
}
|
||||
entry.count += 1;
|
||||
entry.last_seen = now;
|
||||
self.source_expiry
|
||||
.push_back((now + Duration::from_secs(60), addr));
|
||||
true
|
||||
}
|
||||
|
||||
fn expire_sources(&mut self, now: Instant) {
|
||||
while let Some((deadline, addr)) = self.source_expiry.front().copied() {
|
||||
if deadline > now {
|
||||
break;
|
||||
}
|
||||
self.source_expiry.pop_front();
|
||||
if self.sources.get(&addr).is_some_and(|entry| {
|
||||
now.checked_duration_since(entry.last_seen)
|
||||
.unwrap_or_default()
|
||||
>= Duration::from_secs(60)
|
||||
}) {
|
||||
self.sources.remove(&addr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn split_u32_quota(total: u32, workers: usize, worker: usize) -> u32 {
|
||||
let workers = workers.max(1) as u32;
|
||||
total / workers + u32::from((worker as u32) < total % workers)
|
||||
}
|
||||
|
||||
pub(crate) fn split_u64_quota(total: u64, workers: usize, worker: usize) -> u64 {
|
||||
let workers = workers.max(1) as u64;
|
||||
total / workers + u64::from((worker as u64) < total % workers)
|
||||
}
|
||||
|
||||
fn reserve_quota(total: u32) -> u32 {
|
||||
if total == 0 {
|
||||
0
|
||||
} else {
|
||||
total.div_ceil(10).min(total)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn enforces_packet_byte_and_source_limits() {
|
||||
let start = Instant::now();
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
|
||||
let mut limiter = WorkerResponseLimiter::new(2, 200, 1, start);
|
||||
assert_eq!(
|
||||
limiter.acquire(addr, 50, false, start),
|
||||
ResponsePermit::Regular
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(addr, 50, false, start),
|
||||
ResponsePermit::Rejected
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(addr, 50, false, start + Duration::from_secs(1)),
|
||||
ResponsePermit::Regular
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(addr, 513, false, start + Duration::from_secs(2)),
|
||||
ResponsePermit::Rejected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn priority_queries_can_use_the_reserve() {
|
||||
let start = Instant::now();
|
||||
let mut limiter = WorkerResponseLimiter::new(10, 1_000, 100, start);
|
||||
let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
|
||||
let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 2);
|
||||
let priority = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3);
|
||||
assert_eq!(
|
||||
limiter.acquire(first, 50, false, start),
|
||||
ResponsePermit::Regular
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(second, 50, false, start),
|
||||
ResponsePermit::Regular
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(priority, 50, false, start),
|
||||
ResponsePermit::Rejected
|
||||
);
|
||||
assert_eq!(
|
||||
limiter.acquire(priority, 50, true, start),
|
||||
ResponsePermit::PriorityReserve
|
||||
);
|
||||
let exhausted = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4);
|
||||
assert_eq!(
|
||||
limiter.acquire(exhausted, 100, true, start),
|
||||
ResponsePermit::Rejected
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
// 负责提供并发可读的 DHT 路由节点快照和距离查询
|
||||
|
||||
use crate::types::NodeTuple;
|
||||
use rand::seq::IndexedRandom;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct RoutingSnapshot {
|
||||
v4: Vec<NodeTuple>,
|
||||
v6: Vec<NodeTuple>,
|
||||
}
|
||||
|
||||
impl RoutingSnapshot {
|
||||
pub(crate) fn from_nodes(nodes: Vec<NodeTuple>, 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<bool>) -> Vec<NodeTuple> {
|
||||
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<bool>,
|
||||
) -> Vec<NodeTuple> {
|
||||
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<_>>(),
|
||||
vec![2, 3]
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,728 @@
|
||||
// 负责执行有界 BEP-51 采样并将准入后的 infohash 交给 Peer 查找
|
||||
|
||||
use crate::addr::is_valid_node_addr;
|
||||
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::{DiscoverySource, 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,
|
||||
source: DiscoverySource,
|
||||
}
|
||||
|
||||
struct SampleResponse {
|
||||
remote_addr: SocketAddr,
|
||||
tid: TransactionId,
|
||||
response: DhtResponse,
|
||||
}
|
||||
|
||||
struct SampleAdmissionBatch {
|
||||
preferred_node: NodeTuple,
|
||||
source: DiscoverySource,
|
||||
hashes: Vec<[u8; 20]>,
|
||||
}
|
||||
|
||||
pub(crate) type SampleHashAdmissionCallback = Box<
|
||||
dyn Fn(Vec<[u8; 20]>) -> std::pin::Pin<Box<dyn Future<Output = Vec<[u8; 20]>> + Send>>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SampleInfohashesHandle {
|
||||
response_tx: mpsc::Sender<SampleResponse>,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SampleCandidateRouter {
|
||||
sender: mpsc::Sender<NodeTuple>,
|
||||
sample_percent: u8,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
}
|
||||
|
||||
impl SampleCandidateRouter {
|
||||
pub(crate) fn route(&self, node: NodeTuple) -> bool {
|
||||
if self.sample_percent == 0
|
||||
|| !is_valid_node_addr(&node.addr)
|
||||
|| stable_address_bucket(node.addr) >= self.sample_percent
|
||||
{
|
||||
return false;
|
||||
}
|
||||
match self.sender.try_send(node) {
|
||||
Ok(()) => {
|
||||
self.runtime_stats.sample_candidate_routed();
|
||||
self.runtime_stats.set_sample_candidate_queue_depth(
|
||||
self.sender
|
||||
.max_capacity()
|
||||
.saturating_sub(self.sender.capacity()),
|
||||
);
|
||||
true
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||
self.runtime_stats.sample_candidate_fallback();
|
||||
self.runtime_stats
|
||||
.set_sample_candidate_queue_depth(self.sender.max_capacity());
|
||||
false
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Closed(_)) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn sample_candidate_lane(
|
||||
options: &SampleInfohashesOptions,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
) -> (SampleCandidateRouter, mpsc::Receiver<NodeTuple>) {
|
||||
let capacity = options.candidate_queue_capacity.max(1);
|
||||
let (sender, receiver) = mpsc::channel(capacity);
|
||||
runtime_stats.configure_sample_candidate_queue(capacity);
|
||||
(
|
||||
SampleCandidateRouter {
|
||||
sender,
|
||||
sample_percent: options.new_node_sample_percent.min(100),
|
||||
runtime_stats,
|
||||
},
|
||||
receiver,
|
||||
)
|
||||
}
|
||||
|
||||
fn stable_address_bucket(addr: SocketAddr) -> u8 {
|
||||
let mut hash = 0xcbf2_9ce4_8422_2325u64;
|
||||
let mut mix = |byte: u8| {
|
||||
hash ^= u64::from(byte);
|
||||
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
|
||||
};
|
||||
match addr.ip() {
|
||||
std::net::IpAddr::V4(ip) => ip.octets().into_iter().for_each(&mut mix),
|
||||
std::net::IpAddr::V6(ip) => ip.octets().into_iter().for_each(&mut mix),
|
||||
}
|
||||
addr.port().to_be_bytes().into_iter().for_each(&mut mix);
|
||||
(hash % 100) as u8
|
||||
}
|
||||
|
||||
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<ArcSwapOption<SampleHashAdmissionCallback>>,
|
||||
pub(crate) shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
pub(crate) struct SampleInfohashesInputs<'a> {
|
||||
pub(crate) sockets: &'a std::collections::HashMap<SocketAddr, Arc<UdpSocket>>,
|
||||
pub(crate) snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
pub(crate) candidate_rx: mpsc::Receiver<NodeTuple>,
|
||||
pub(crate) crawl_engine: Arc<CrawlEngine>,
|
||||
pub(crate) peer_lookup: PeerLookupHandle,
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_sample_infohashes(
|
||||
netmode: NetMode,
|
||||
local_id: [u8; 20],
|
||||
inputs: SampleInfohashesInputs<'_>,
|
||||
runtime: SampleInfohashesRuntime,
|
||||
) -> SampleInfohashesHandle {
|
||||
let SampleInfohashesInputs {
|
||||
sockets,
|
||||
snapshot,
|
||||
candidate_rx,
|
||||
crawl_engine,
|
||||
peer_lookup,
|
||||
} = inputs;
|
||||
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(),
|
||||
options.fallback_to_iterative,
|
||||
stats.clone(),
|
||||
shutdown.clone(),
|
||||
);
|
||||
let actor = SampleInfohashesActor {
|
||||
netmode,
|
||||
local_id,
|
||||
socket_v4,
|
||||
socket_v6,
|
||||
snapshot,
|
||||
candidate_rx,
|
||||
direct_only: options.new_node_sample_percent > 0,
|
||||
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<Arc<UdpSocket>>,
|
||||
socket_v6: Option<Arc<UdpSocket>>,
|
||||
snapshot: Arc<ArcSwap<RoutingSnapshot>>,
|
||||
candidate_rx: mpsc::Receiver<NodeTuple>,
|
||||
direct_only: bool,
|
||||
crawl_engine: Arc<CrawlEngine>,
|
||||
admission_tx: mpsc::Sender<SampleAdmissionBatch>,
|
||||
response_rx: mpsc::Receiver<SampleResponse>,
|
||||
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<SocketAddr, Instant>,
|
||||
pending_addrs: AHashSet<SocketAddr>,
|
||||
pending: AHashMap<PendingKey, PendingRequest>,
|
||||
pending_expiry: VecDeque<(Instant, PendingKey)>,
|
||||
next_tid: u64,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
outbound_query_budget: SharedRateBudget,
|
||||
shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
fn spawn_sample_admission(
|
||||
mut receiver: mpsc::Receiver<SampleAdmissionBatch>,
|
||||
hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
|
||||
peer_lookup_tx: mpsc::Sender<PeerLookupRequest>,
|
||||
fallback_to_iterative: bool,
|
||||
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),
|
||||
allow_iterative_fallback: fallback_to_iterative,
|
||||
source: batch.source,
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.runtime_stats.set_sample_candidate_queue_depth(0);
|
||||
}
|
||||
|
||||
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;
|
||||
while sent_count < budget {
|
||||
let Some(node) = self.next_direct_candidate(now) else {
|
||||
break;
|
||||
};
|
||||
if self
|
||||
.send_query(node, DiscoverySource::SampleDirect, now)
|
||||
.await
|
||||
{
|
||||
sent_count += 1;
|
||||
}
|
||||
}
|
||||
if self.direct_only {
|
||||
if sent_count < budget {
|
||||
self.query_budget.refund(budget - sent_count);
|
||||
}
|
||||
return;
|
||||
}
|
||||
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, DiscoverySource::SampleSnapshot, now)
|
||||
.await
|
||||
{
|
||||
sent_count += 1;
|
||||
}
|
||||
}
|
||||
if sent_count < budget {
|
||||
self.query_budget.refund(budget - sent_count);
|
||||
}
|
||||
}
|
||||
|
||||
fn next_direct_candidate(&mut self, now: Instant) -> Option<NodeTuple> {
|
||||
for _ in 0..64 {
|
||||
let node = self.candidate_rx.try_recv().ok()?;
|
||||
self.runtime_stats
|
||||
.set_sample_candidate_queue_depth(self.candidate_rx.len());
|
||||
if self.pending_addrs.contains(&node.addr)
|
||||
|| self
|
||||
.next_allowed
|
||||
.get(&node.addr)
|
||||
.is_some_and(|deadline| *deadline > now)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
return Some(node);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn send_query(&mut self, node: NodeTuple, source: DiscoverySource, 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, source });
|
||||
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(source == DiscoverySource::SampleDirect);
|
||||
#[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,
|
||||
};
|
||||
let Some(pending) = self.pending.remove(&key) else {
|
||||
return;
|
||||
};
|
||||
self.pending_addrs.remove(&event.remote_addr);
|
||||
self.runtime_stats
|
||||
.sample_response(pending.source == DiscoverySource::SampleDirect);
|
||||
|
||||
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,
|
||||
source: pending.source,
|
||||
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;
|
||||
|
||||
fn node_for_lane(want_direct: bool, percent: u8) -> NodeTuple {
|
||||
(1..=u16::MAX)
|
||||
.map(|port| NodeTuple {
|
||||
id: [port as u8; 20],
|
||||
addr: SocketAddr::from(([8, 8, 8, 8], port)),
|
||||
})
|
||||
.find(|node| (stable_address_bucket(node.addr) < percent) == want_direct)
|
||||
.expect("both routing lanes have an address")
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_node_routing_is_stable_and_bounded() {
|
||||
let options = SampleInfohashesOptions {
|
||||
new_node_sample_percent: 50,
|
||||
candidate_queue_capacity: 1,
|
||||
..SampleInfohashesOptions::default()
|
||||
};
|
||||
let stats = DhtRuntimeStats::default();
|
||||
let (router, mut receiver) = sample_candidate_lane(&options, stats.clone());
|
||||
let direct = node_for_lane(true, 50);
|
||||
let crawl = node_for_lane(false, 50);
|
||||
|
||||
assert!(router.route(direct));
|
||||
assert!(!router.route(crawl));
|
||||
assert!(!router.route(direct));
|
||||
assert_eq!(receiver.try_recv().unwrap(), direct);
|
||||
|
||||
let snapshot = stats.snapshot();
|
||||
assert_eq!(snapshot.sample_candidate_queue_capacity, 1);
|
||||
assert_eq!(snapshot.sample_candidates_routed, 1);
|
||||
assert_eq!(snapshot.sample_candidates_fallback, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_percent_preserves_snapshot_only_sampling() {
|
||||
let options = SampleInfohashesOptions {
|
||||
new_node_sample_percent: 0,
|
||||
..SampleInfohashesOptions::default()
|
||||
};
|
||||
let (router, mut receiver) = sample_candidate_lane(&options, DhtRuntimeStats::default());
|
||||
assert!(!router.route(node_for_lane(false, 0)));
|
||||
assert!(receiver.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn batch_admission_only_forwards_application_approved_hashes() {
|
||||
let (batch_tx, batch_rx) = mpsc::channel(1);
|
||||
let (lookup_tx, mut lookup_rx) = mpsc::channel(2);
|
||||
let admission = Arc::new(ArcSwapOption::empty());
|
||||
let callback: Arc<SampleHashAdmissionCallback> = Arc::new(Box::new(|hashes| {
|
||||
Box::pin(async move { hashes.into_iter().filter(|hash| *hash == [2; 20]).collect() })
|
||||
}));
|
||||
admission.store(Some(callback));
|
||||
let stats = DhtRuntimeStats::default();
|
||||
let shutdown = CancellationToken::new();
|
||||
spawn_sample_admission(
|
||||
batch_rx,
|
||||
admission,
|
||||
lookup_tx,
|
||||
false,
|
||||
stats.clone(),
|
||||
shutdown.clone(),
|
||||
);
|
||||
batch_tx
|
||||
.send(SampleAdmissionBatch {
|
||||
preferred_node: NodeTuple {
|
||||
id: [9; 20],
|
||||
addr: "127.0.0.1:6881".parse().unwrap(),
|
||||
},
|
||||
source: DiscoverySource::SampleDirect,
|
||||
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]);
|
||||
assert!(!request.allow_iterative_fallback);
|
||||
assert_eq!(request.source, DiscoverySource::SampleDirect);
|
||||
let snapshot = stats.snapshot();
|
||||
assert_eq!(snapshot.sample_infohashes_hashes_filtered, 1);
|
||||
assert_eq!(snapshot.sample_infohashes_hashes_discovered, 1);
|
||||
shutdown.cancel();
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,808 @@
|
||||
// 负责组合 DHT 网络组件回调生命周期和公开运行接口
|
||||
|
||||
use crate::addr::is_valid_node_addr;
|
||||
use crate::budget::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::response_limiter::{
|
||||
ResponsePermit, WorkerResponseLimiter, split_u32_quota, split_u64_quota,
|
||||
};
|
||||
use crate::runtime_stats::{DhtRuntimeLimits, DhtRuntimeStats};
|
||||
use crate::sample_infohashes::{
|
||||
SampleHashAdmissionCallback, SampleInfohashesHandle, SampleInfohashesInputs,
|
||||
SampleInfohashesRuntime, is_sample_infohashes_tid, sample_candidate_lane,
|
||||
spawn_sample_infohashes,
|
||||
};
|
||||
use crate::scheduler::{
|
||||
MetadataCompletionCallback, MetadataFetchCallback, MetadataScheduler,
|
||||
MetadataSchedulerCallbacks, MetadataSchedulerLimits, MetadataSchedulerRuntime,
|
||||
TorrentAckCallback,
|
||||
};
|
||||
use crate::types::{
|
||||
DHTOptions, DiscoverySource, HashDiscovered, MetadataFetchCompletion, NetMode, NodeTuple,
|
||||
TorrentInfo,
|
||||
};
|
||||
use crate::udp_buffer::UdpBufferPool;
|
||||
use crate::udp_ingress::{WorkerHandle, spawn_udp_listener};
|
||||
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;
|
||||
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<dyn Fn(&str) -> bool + Send + Sync + 'static>;
|
||||
type ErrorCallback = Box<dyn Fn(crate::error::DHTError) + Send + Sync + 'static>;
|
||||
|
||||
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)]
|
||||
/// 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<HashMap<SocketAddr, Arc<UdpSocket>>>,
|
||||
token_secret: [u8; 10],
|
||||
torrent_callback: Arc<ArcSwapOption<TorrentAckCallback>>,
|
||||
hash_filter: Arc<ArcSwapOption<FilterCallback>>,
|
||||
on_metadata_fetch: Arc<ArcSwapOption<MetadataFetchCallback>>,
|
||||
sample_hash_admission: Arc<ArcSwapOption<SampleHashAdmissionCallback>>,
|
||||
metadata_completion_callback: Arc<ArcSwapOption<MetadataCompletionCallback>>,
|
||||
error_callback: Arc<ArcSwapOption<ErrorCallback>>,
|
||||
crawl_engine: Arc<CrawlEngine>,
|
||||
peer_lookup: PeerLookupHandle,
|
||||
metadata_fetcher: Arc<RbitFetcher>,
|
||||
sample_infohashes: SampleInfohashesHandle,
|
||||
hash_events_tx: mpsc::Sender<HashDiscovered>,
|
||||
metadata_queue_len: Arc<AtomicUsize>,
|
||||
max_metadata_queue_size: usize,
|
||||
runtime_stats: DhtRuntimeStats,
|
||||
shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
fn create_udp_sock(domain: Domain, ty: Type, addr: SocketAddr) -> std::io::Result<UdpSocket> {
|
||||
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())
|
||||
}
|
||||
|
||||
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<Self> {
|
||||
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::<HashDiscovered>(options.hash_queue_capacity);
|
||||
let fetcher = Arc::new(RbitFetcher::new_with_runtime_stats(
|
||||
options.metadata.timeout_secs,
|
||||
options.metadata.max_connects_per_second,
|
||||
options.metadata.max_metadata_size_bytes,
|
||||
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 (sample_candidate_router, sample_candidate_rx) =
|
||||
sample_candidate_lane(&options.sample_infohashes, runtime_stats.clone());
|
||||
let crawl_engine = Arc::new(CrawlEngine::new(
|
||||
crawl_config.clone(),
|
||||
runtime_stats.clone(),
|
||||
outbound_query_budget.clone(),
|
||||
Some(sample_candidate_router),
|
||||
));
|
||||
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,
|
||||
SampleInfohashesInputs {
|
||||
sockets: &sockets_by_bind_addr,
|
||||
snapshot: crawl_engine.snapshot.clone(),
|
||||
candidate_rx: sample_candidate_rx,
|
||||
crawl_engine: crawl_engine.clone(),
|
||||
peer_lookup: 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<F, Fut>(&self, callback: F)
|
||||
where
|
||||
F: Fn(String) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = bool> + Send + 'static,
|
||||
{
|
||||
let callback: Arc<MetadataFetchCallback> =
|
||||
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<F, Fut>(&self, callback: F)
|
||||
where
|
||||
F: Fn(Vec<[u8; 20]>) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = Vec<[u8; 20]>> + Send + 'static,
|
||||
{
|
||||
let callback: Arc<SampleHashAdmissionCallback> =
|
||||
Arc::new(Box::new(move |hashes| Box::pin(callback(hashes))));
|
||||
self.sample_hash_admission.store(Some(callback));
|
||||
}
|
||||
|
||||
/// Registers a torrent callback whose return is implicitly treated as accepted delivery.
|
||||
///
|
||||
/// Registering a new torrent callback replaces the previous one.
|
||||
pub fn on_torrent<F>(&self, callback: F)
|
||||
where
|
||||
F: Fn(TorrentInfo) + Send + Sync + 'static,
|
||||
{
|
||||
let callback: Arc<TorrentAckCallback> = 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<F>(&self, callback: F)
|
||||
where
|
||||
F: Fn(TorrentInfo) -> bool + Send + Sync + 'static,
|
||||
{
|
||||
let callback: Arc<TorrentAckCallback> = 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<F>(&self, callback: F)
|
||||
where
|
||||
F: Fn(MetadataFetchCompletion) + Send + Sync + 'static,
|
||||
{
|
||||
let callback: Arc<MetadataCompletionCallback> = 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<F>(&self, filter: F)
|
||||
where
|
||||
F: Fn(&str) -> bool + Send + Sync + 'static,
|
||||
{
|
||||
let filter: Arc<FilterCallback> = Arc::new(Box::new(filter));
|
||||
self.hash_filter.store(Some(filter));
|
||||
}
|
||||
|
||||
/// Alias for [`Self::filter`].
|
||||
pub fn set_filter<F>(&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<F>(&self, callback: F)
|
||||
where
|
||||
F: Fn(crate::error::DHTError) + Send + Sync + 'static,
|
||||
{
|
||||
let callback: Arc<ErrorCallback> = 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<crate::PeerLookupResult> {
|
||||
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<WorkerHandle> {
|
||||
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<WorkerHandle> = 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),
|
||||
source: DiscoverySource::AnnouncePeer,
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
// 负责定义 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<FileInfo>,
|
||||
/// Torrent piece length in bytes, or zero if absent.
|
||||
pub piece_length: u64,
|
||||
/// Peer addresses used to obtain the Metadata.
|
||||
pub peers: Vec<String>,
|
||||
/// 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 accepted BEP-9 info dictionary size in bytes.
|
||||
pub max_metadata_size_bytes: usize,
|
||||
/// 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, Copy, PartialEq, Eq)]
|
||||
/// Network path that supplied a Peer candidate for Metadata download.
|
||||
pub enum DiscoverySource {
|
||||
/// A remote node announced itself as a Peer.
|
||||
AnnouncePeer,
|
||||
/// A Peer was found for a Hash sampled directly from a newly discovered node.
|
||||
SampleDirect,
|
||||
/// A Peer was found for a Hash sampled from the responsive-node snapshot.
|
||||
SampleSnapshot,
|
||||
/// A Peer was found by the Metadata scheduler's delayed active lookup.
|
||||
ActiveLookup,
|
||||
}
|
||||
|
||||
#[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,
|
||||
/// Network path that supplied this Peer candidate.
|
||||
pub source: DiscoverySource,
|
||||
/// Monotonic discovery time used for freshness and queue ordering.
|
||||
pub discovered_at: std::time::Instant,
|
||||
}
|
||||
|
||||
#[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<SocketAddr>,
|
||||
/// 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,
|
||||
/// Percentage of newly discovered node addresses routed directly to BEP-51 sampling.
|
||||
pub new_node_sample_percent: u8,
|
||||
/// Capacity of the bounded newly discovered node sampling lane.
|
||||
pub candidate_queue_capacity: usize,
|
||||
/// Whether a sampled Hash should fall back to iterative get_peers after the sampling node
|
||||
/// returns no Peer.
|
||||
pub fallback_to_iterative: bool,
|
||||
/// 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<String>,
|
||||
/// 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,
|
||||
max_metadata_size_bytes: 10 * 1024 * 1024,
|
||||
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,
|
||||
new_node_sample_percent: 0,
|
||||
candidate_queue_capacity: 4_096,
|
||||
fallback_to_iterative: true,
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
// 负责复用有界 UDP 接收缓冲区并控制内存占用
|
||||
|
||||
//! 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<PoolInner>,
|
||||
}
|
||||
|
||||
struct PoolInner {
|
||||
free: ArrayQueue<Box<[u8]>>,
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,347 @@
|
||||
// 负责将 UDP 数据包分发到有界工作队列并处理过载
|
||||
|
||||
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<UdpSocket>,
|
||||
mut workers: Vec<WorkerHandle>,
|
||||
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<WorkerHandle>,
|
||||
) -> 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user