feat: 重组项目结构并完善运行管理

This commit is contained in:
chuan
2026-08-10 15:00:39 +08:00
parent ea4625e5da
commit 28fa69e614
127 changed files with 3927 additions and 1079 deletions
+49
View File
@@ -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% 保底预算。
+54
View File
@@ -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"
+272
View File
@@ -0,0 +1,272 @@
# dht-crawler
[![Crates.io](https://img.shields.io/crates/v/dht-crawler.svg)](https://crates.io/crates/dht-crawler)
[![Documentation](https://docs.rs/dht-crawler/badge.svg)](https://docs.rs/dht-crawler)
[![License](https://img.shields.io/crates/l/dht-crawler.svg)](../../LICENSE)
基于 Rust 和 Tokio 的 BitTorrent DHT 爬虫库。它参与 BEP-5 DHT 网络,通过 BEP-51
`sample_infohashes` 主动发现 InfoHash,也接收 `announce_peer`,并通过 BEP-9
`ut_metadata` 获取、校验和解析 torrent 元数据。
`dht-crawler` 提供:
- IPv4、IPv6 和双栈 DHT
- 主动节点发现、BEP-51 InfoHash 采样与 `get_peers` 查询;
- 有界、去重的 Metadata 下载队列;
- InfoHash 过滤、异步准入、结果交付和完成通知;
- 默认可用的运行时统计,以及可选的 `metrics` 集成。
## 安装
```bash
cargo add dht-crawler
cargo add tokio --features rt-multi-thread,macros,signal
```
或在 `Cargo.toml` 中添加:
```toml
[dependencies]
dht-crawler = "0.2"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] }
```
## 快速开始
```rust
use dht_crawler::prelude::*;
#[tokio::main]
async fn main() -> Result<()> {
let server = DHTServer::new(DHTOptions {
port: 6881,
netmode: NetMode::Ipv4Only,
..Default::default()
})
.await?;
server.on_torrent(|torrent| {
println!(
"{} {} {}",
torrent.info_hash,
torrent.name,
torrent.format_size()
);
});
server.on_error(|error| {
eprintln!("DHT runtime error: {error}");
});
let shutdown = server.clone();
tokio::spawn(async move {
if tokio::signal::ctrl_c().await.is_ok() {
shutdown.shutdown();
}
});
// 一直运行,直到 shutdown() 被调用。
server.start().await
}
```
运行仓库中的完整示例:
```bash
cargo run --release --example dht_crawler_example
```
## 核心 API
`DHTServer` 是主要入口:
| API | 用途 |
|---|---|
| `DHTServer::new(options)` | 校验配置、绑定 UDP Socket 并创建内部管道 |
| `start().await` | 启动 DHT 与爬取任务,等待 `shutdown()` |
| `shutdown()` | 停止 UDP、爬取和 Metadata 任务;可重复调用 |
| `filter(callback)` | 在 InfoHash 进入队列前执行同步过滤 |
| `on_metadata_fetch(callback)` | 在第一次 Peer 下载前执行异步准入 |
| `on_torrent(callback)` | 接收已校验的 `TorrentInfo` |
| `on_torrent_with_ack(callback)` | 接收结果并显式确认是否接受交付 |
| `on_metadata_fetch_complete(callback)` | 接收已准入任务的最终状态 |
| `on_error(callback)` | 接收运行期错误 |
| `runtime_stats()` | 获取可复制的运行时统计句柄 |
同类回调重复注册时,新回调会替换旧回调。
### 过滤与准入
`filter` 是同步的早期过滤器,适合拦截已处理过的 InfoHash:
```rust
server.filter(|info_hash| !already_exists(info_hash));
```
`on_metadata_fetch` 是异步准入回调,在实际连接 Peer 前调用:
```rust
server.on_metadata_fetch(|info_hash| async move {
should_download(&info_hash).await
});
```
返回 `false` 会终止任务,不下载 Metadata,也不会触发 torrent 或 completion 回调。
未注册准入回调时默认允许下载。
### 交付确认
不需要确认下游是否接收时使用 `on_torrent`。需要确认下游是否成功接收时使用
`on_torrent_with_ack`
```rust
server.on_torrent_with_ack(|torrent| {
output.try_send(torrent).is_ok()
});
server.on_metadata_fetch_complete(|completion| {
println!(
"{}: {:?}, attempts={}",
completion.info_hash,
completion.status,
completion.attempts
);
});
```
完成状态:
| 状态 | 含义 |
|---|---|
| `Accepted` | Metadata 下载成功,结果已被回调接受 |
| `FetchFailed` | 所有可用 Peer 尝试均失败 |
| `DeliveryRejected` | Metadata 下载成功,但结果未被回调接受 |
`attempts` 只统计实际发起的 Peer 网络请求。异步准入拒绝不会产生 completion 事件。
## 配置
大多数调用方可以从 `DHTOptions::default()` 开始,只覆盖监听方式和容量限制:
```rust
let options = DHTOptions {
port: 6881,
netmode: NetMode::DualStack,
hash_queue_capacity: 20_000,
metadata: MetadataOptions {
timeout_secs: 5,
max_queue_size: 20_000,
max_worker_count: 8,
max_connects_per_second: 2,
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)
+110
View File
@@ -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 延迟。
+151
View File
@@ -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(())
}
+73
View File
@@ -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()));
}
}
+223
View File
@@ -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));
}
}
+182
View File
@@ -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));
}
}
+139
View File
@@ -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);
}
}
+974
View File
@@ -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);
}
}
+30
View File
@@ -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>;
+310
View File
@@ -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());
}
}
+63
View File
@@ -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,
};
}
+740
View File
@@ -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();
}
}
+84
View File
@@ -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]);
}
}
+392
View File
@@ -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());
}
}
+911
View File
@@ -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();
}
}
+61
View File
@@ -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>,
}
+268
View File
@@ -0,0 +1,268 @@
// 负责限制 DHT 响应的包速率字节速率单来源速率和优先保留预算
use std::{
collections::VecDeque,
net::SocketAddr,
time::{Duration, Instant},
};
use ahash::AHashMap;
use crate::budget::RateBucket;
#[derive(Clone, Copy)]
struct SourceResponseWindow {
started_at: Instant,
last_seen: Instant,
count: u32,
}
pub(crate) struct WorkerResponseLimiter {
regular_packets: RateBucket,
regular_bytes: RateBucket,
priority_packets: RateBucket,
priority_bytes: RateBucket,
per_source_rate: u32,
sources: AHashMap<SocketAddr, SourceResponseWindow>,
source_expiry: VecDeque<(Instant, SocketAddr)>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ResponsePermit {
Regular,
PriorityReserve,
Rejected,
}
impl WorkerResponseLimiter {
pub(crate) fn new(
packet_rate: u32,
byte_rate: u64,
per_source_rate: u32,
now: Instant,
) -> Self {
let byte_rate = byte_rate.min(u32::MAX as u64) as u32;
let priority_packet_rate = reserve_quota(packet_rate);
let priority_byte_rate = reserve_quota(byte_rate);
let packet_rate = packet_rate.saturating_sub(priority_packet_rate);
let byte_rate = byte_rate.saturating_sub(priority_byte_rate);
Self {
regular_packets: RateBucket::per_second(
packet_rate,
packet_rate.div_ceil(5).max(1),
true,
now,
),
regular_bytes: RateBucket::per_second(
byte_rate,
byte_rate.div_ceil(5).max(512),
true,
now,
),
priority_packets: RateBucket::per_second(
priority_packet_rate,
priority_packet_rate.div_ceil(5).max(1),
true,
now,
),
priority_bytes: RateBucket::per_second(
priority_byte_rate,
priority_byte_rate.div_ceil(5).max(512),
true,
now,
),
per_source_rate,
sources: AHashMap::new(),
source_expiry: VecDeque::new(),
}
}
pub(crate) fn acquire(
&mut self,
addr: SocketAddr,
encoded_len: usize,
is_priority: bool,
now: Instant,
) -> ResponsePermit {
self.expire_sources(now);
if Self::take_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
now,
) {
if self.acquire_source_slot(addr, now) {
return ResponsePermit::Regular;
}
Self::refund_budget(
&mut self.regular_packets,
&mut self.regular_bytes,
encoded_len,
);
return ResponsePermit::Rejected;
}
if is_priority
&& Self::take_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
now,
)
{
if self.acquire_source_slot(addr, now) {
return ResponsePermit::PriorityReserve;
}
Self::refund_budget(
&mut self.priority_packets,
&mut self.priority_bytes,
encoded_len,
);
}
ResponsePermit::Rejected
}
fn take_budget(
packets: &mut RateBucket,
bytes: &mut RateBucket,
encoded_len: usize,
now: Instant,
) -> bool {
if !packets.try_take_one(now) {
return false;
}
if !bytes.try_take_exact(encoded_len, now) {
packets.refund_one();
return false;
}
true
}
fn refund_budget(packets: &mut RateBucket, bytes: &mut RateBucket, encoded_len: usize) {
packets.refund_one();
bytes.refund(encoded_len);
}
fn acquire_source_slot(&mut self, addr: SocketAddr, now: Instant) -> bool {
if self.per_source_rate == 0 {
return false;
}
let entry = self.sources.entry(addr).or_insert(SourceResponseWindow {
started_at: now,
last_seen: now,
count: 0,
});
if now
.checked_duration_since(entry.started_at)
.unwrap_or_default()
>= Duration::from_secs(1)
{
entry.started_at = now;
entry.count = 0;
}
if entry.count >= self.per_source_rate {
return false;
}
entry.count += 1;
entry.last_seen = now;
self.source_expiry
.push_back((now + Duration::from_secs(60), addr));
true
}
fn expire_sources(&mut self, now: Instant) {
while let Some((deadline, addr)) = self.source_expiry.front().copied() {
if deadline > now {
break;
}
self.source_expiry.pop_front();
if self.sources.get(&addr).is_some_and(|entry| {
now.checked_duration_since(entry.last_seen)
.unwrap_or_default()
>= Duration::from_secs(60)
}) {
self.sources.remove(&addr);
}
}
}
}
pub(crate) fn split_u32_quota(total: u32, workers: usize, worker: usize) -> u32 {
let workers = workers.max(1) as u32;
total / workers + u32::from((worker as u32) < total % workers)
}
pub(crate) fn split_u64_quota(total: u64, workers: usize, worker: usize) -> u64 {
let workers = workers.max(1) as u64;
total / workers + u64::from((worker as u64) < total % workers)
}
fn reserve_quota(total: u32) -> u32 {
if total == 0 {
0
} else {
total.div_ceil(10).min(total)
}
}
#[cfg(test)]
mod tests {
use std::{
net::{IpAddr, Ipv4Addr},
time::Duration,
};
use super::*;
#[test]
fn enforces_packet_byte_and_source_limits() {
let start = Instant::now();
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
let mut limiter = WorkerResponseLimiter::new(2, 200, 1, start);
assert_eq!(
limiter.acquire(addr, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(addr, 50, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(addr, 50, false, start + Duration::from_secs(1)),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(addr, 513, false, start + Duration::from_secs(2)),
ResponsePermit::Rejected
);
}
#[test]
fn priority_queries_can_use_the_reserve() {
let start = Instant::now();
let mut limiter = WorkerResponseLimiter::new(10, 1_000, 100, start);
let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1);
let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 2);
let priority = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 3);
assert_eq!(
limiter.acquire(first, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(second, 50, false, start),
ResponsePermit::Regular
);
assert_eq!(
limiter.acquire(priority, 50, false, start),
ResponsePermit::Rejected
);
assert_eq!(
limiter.acquire(priority, 50, true, start),
ResponsePermit::PriorityReserve
);
let exhausted = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4);
assert_eq!(
limiter.acquire(exhausted, 100, true, start),
ResponsePermit::Rejected
);
}
}
+109
View File
@@ -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
+728
View File
@@ -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
+808
View File
@@ -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
}
}
+447
View File
@@ -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,
}
}
}
+81
View File
@@ -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()
}
+347
View File
@@ -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);
}
}