tools/dht
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Prepare for crates.io release v0.0.1

Browse Source
This commit is contained in:
桥下红药
2026-01-16 19:46:51 +08:00
parent cb6f005967
commit c356926d69
10 changed files with 136 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+3 -3
View File
@@ -2,9 +2,9 @@ mod error;
mod server;
pub mod protocol;
pub mod types;
pub mod metadata; // 公开 metadata 模块
mod sharded; // 分片锁模块
pub mod scheduler; // 元数据调度器
pub mod metadata;
mod sharded;
pub mod scheduler;
pub use error::{DHTError, Result};
pub use server::{DHTServer, HashDiscovered};
src/main.rs
-106
View File
@@ -1,106 +0,0 @@
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;
use dht_crawler::prelude::*;
use std::sync::Arc;
use mimalloc::MiMalloc;
use tracing_subscriber::EnvFilter;
use std::sync::atomic::{AtomicUsize, Ordering};
#[tokio::main]
async fn main() -> Result<()> {
if std::env::var("RUST_LOG").is_err() {
std::env::set_var("RUST_LOG", "info");
}
// 直接输出到 stdout，避免 _guard 被 drop 导致日志丢失
tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env())
.with_ansi(true)
.init();
let options = DHTOptions {
port: 12313,
auto_metadata: true,
metadata_timeout: 3, // ✅ 快速超时，快速失败
max_metadata_queue_size: 100000, // ✅ 大缓冲区（防止饱和）
max_metadata_worker_count: 1000, // ✅ 激进并发（最大化吞吐）
netmode: NetMode::DualStack, // 网络模式：Ipv4Only（仅IPv4）、Ipv6Only（仅IPv6）、DualStack（双栈，默认）
};
// 统计计数器
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);
// 设置 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
});
server.set_filter(|_hash| {
true
});
server.on_duplicate(|_hash| {
});
// 启动监控任务
let dht_monitor = server.clone();
let count_monitor = torrent_count.clone();
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();
// ✅ 监控：布隆过滤器的位使用情况反映了爬虫的活跃度
log::info!(
"📊 [监控] 时长: {}s | 成功抓取: ✨ {} | 活跃指纹: {}",
uptime, success_fetch, dht_monitor.get_seen_count()
);
if uptime > 0 && success_fetch > 0 {
let speed = (success_fetch as f64) / (uptime as f64 / 60.0);
log::info!("📈 平均抓取速度: {:.2} 种子/分钟", speed);
}
}
});
server.start().await?;
Ok(())
}
src/metadata.rs
-1
View File
@@ -33,7 +33,6 @@ impl RbitFetcher {
let peer_id = PeerId::generate();
// 🔥 修改点：缩短连接超时到 3 秒
// DHT 网络很不稳定，如果 3 秒连不上，基本就是连不上了，不要浪费时间
let mut conn = match timeout(
Duration::from_secs(3),
src/scheduler.rs
+50 -22
View File
@@ -4,6 +4,7 @@ use crate::metadata::RbitFetcher;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use tokio::sync::{mpsc, Mutex};
#[cfg(debug_assertions)]
use std::time::Duration;
type TorrentCallback = Arc<dyn Fn(TorrentInfo) + Send + Sync>;
@@ -124,37 +125,63 @@ impl MetadataScheduler {
}
// 主循环：只负责接收 hash 并转发到 worker 队列
#[cfg(debug_assertions)]
let mut stats_interval = tokio::time::interval(Duration::from_secs(60));
#[cfg(debug_assertions)]
stats_interval.tick().await;
loop {
tokio::select! {
Some(hash) = self.hash_rx.recv() => {
self.total_received.fetch_add(1, Ordering::Relaxed);
// 尝试发送到 worker 队列
match task_tx.try_send(hash) {
Ok(_) => {
// 成功入队，增加计数器
self.queue_len.fetch_add(1, Ordering::Relaxed);
#[cfg(debug_assertions)]
{
tokio::select! {
Some(hash) = self.hash_rx.recv() => {
self.total_received.fetch_add(1, Ordering::Relaxed);
// 尝试发送到 worker 队列
match task_tx.try_send(hash) {
Ok(_) => {
// 成功入队，增加计数器
self.queue_len.fetch_add(1, Ordering::Relaxed);
}
Err(mpsc::error::TrySendError::Full(_)) => {
// 队列满，丢弃
self.total_dropped.fetch_add(1, Ordering::Relaxed);
}
Err(_) => break, // Channel 关闭
}
Err(mpsc::error::TrySendError::Full(_)) => {
// 队列满，丢弃
self.total_dropped.fetch_add(1, Ordering::Relaxed);
}
Err(_) => break, // Channel 关闭
}
_ = stats_interval.tick() => {
self.print_stats(&task_tx);
}
else => break,
}
_ = stats_interval.tick() => {
self.print_stats(&task_tx);
}
#[cfg(not(debug_assertions))]
{
match self.hash_rx.recv().await {
Some(hash) => {
self.total_received.fetch_add(1, Ordering::Relaxed);
// 尝试发送到 worker 队列
match task_tx.try_send(hash) {
Ok(_) => {
// 成功入队，增加计数器
self.queue_len.fetch_add(1, Ordering::Relaxed);
}
Err(mpsc::error::TrySendError::Full(_)) => {
// 队列满，丢弃
self.total_dropped.fetch_add(1, Ordering::Relaxed);
}
Err(_) => break, // Channel 关闭
}
}
None => break, // Channel 关闭
}
else => break,
}
}
// log::info!("🛑 Metadata 调度器停止");
}
/// 处理单个 hash（Worker 调用）
@@ -221,7 +248,8 @@ impl MetadataScheduler {
}
}
/// 输出统计信息
/// 输出统计信息（仅在 debug 模式下编译）
#[cfg(debug_assertions)]
fn print_stats(&self, task_tx: &mpsc::Sender<HashDiscovered>) {
let received = self.total_received.load(Ordering::Relaxed);
let dropped = self.total_dropped.load(Ordering::Relaxed);
src/server.rs
+4 -43
View File
@@ -3,7 +3,7 @@ use crate::metadata::RbitFetcher;
use crate::protocol::{DhtMessage, DhtArgs, DhtResponse};
use crate::scheduler::MetadataScheduler;
use crate::types::{DHTOptions, TorrentInfo, NetMode};
use crate::sharded::{ShardedBloom, ShardedNodeQueue, NodeTuple};
use crate::sharded::{ShardedNodeQueue, NodeTuple};
use rand::Rng;
use ahash::AHasher;
use std::hash::{Hash, Hasher};
@@ -40,7 +40,6 @@ pub struct HashDiscovered {
type TorrentCallback = Arc<dyn Fn(TorrentInfo) + Send + Sync>;
type FilterCallback = Arc<dyn Fn(&str) -> bool + Send + Sync>;
type DuplicateCallback = Arc<dyn Fn(&str) + Send + Sync>;
#[derive(Clone)]
pub struct DHTServer {
@@ -53,12 +52,10 @@ pub struct DHTServer {
callback: Arc<RwLock<Option<TorrentCallback>>>,
filter: Arc<RwLock<Option<FilterCallback>>>,
on_duplicate: Arc<RwLock<Option<DuplicateCallback>>>,
on_metadata_fetch: Arc<RwLock<Option<MetadataFetchCallback>>>,
// 使用分片锁，大幅减少竞争
node_queue: Arc<ShardedNodeQueue>,
seen_hashes: Arc<ShardedBloom>,
// 发送 hash 发现事件
hash_tx: mpsc::Sender<HashDiscovered>,
@@ -138,13 +135,9 @@ impl DHTServer {
let mut rng = rand::thread_rng();
let token_secret: Vec<u8> = (0..10).map(|_| rng.gen()).collect();
// 使用分片队列和分片布隆过滤器
// 队列容量：100000 个节点（扩容以适应 DHT 网络裂变速度）
let node_queue = ShardedNodeQueue::new(100000);
// 布隆过滤器：预期500万元素，0.1%误判率
// 内存使用：约 90MB（32分片 × 2.8MB）
let bloom = ShardedBloom::new_for_fp_rate(5_000_000, 0.001);
// 使用分片队列
// 队列容量：从配置获取
let node_queue = ShardedNodeQueue::new(options.node_queue_capacity);
// -----------------------------------------------------------
// 内部初始化 MetadataScheduler
@@ -184,9 +177,7 @@ impl DHTServer {
callback,
on_metadata_fetch,
node_queue: Arc::new(node_queue),
seen_hashes: Arc::new(bloom),
filter: Arc::new(RwLock::new(None)),
on_duplicate: Arc::new(RwLock::new(None)),
hash_tx,
metadata_queue_len,
max_metadata_queue_size: options.max_metadata_queue_size,
@@ -302,22 +293,6 @@ impl DHTServer {
*self.filter.write().unwrap() = Some(Arc::new(filter));
}
/// 设置重复 Hash 发现的回调
///
/// 当接收到的 Hash 已经被布隆过滤器标记为“已存在”时调用。
///
/// # 注意事项
/// - 库内部已经自动为每次调用包裹了 `tokio::spawn`。
/// - 因此你可以放心地在回调中执行耗时操作（如数据库记录），而不用担心阻塞 UDP 线程。
/// - 虽然内部有 spawn，但频繁触发仍会产生大量任务，请注意资源控制。
pub fn on_duplicate<F>(&self, callback: F) where F: Fn(&str) + Send + Sync + 'static {
*self.on_duplicate.write().unwrap() = Some(Arc::new(callback));
}
pub fn get_seen_count(&self) -> usize {
// 分片布隆过滤器的位数统计
self.seen_hashes.number_of_bits() as usize
}
pub fn get_node_pool_size(&self) -> usize {
self.node_queue.len()
@@ -566,20 +541,6 @@ impl DHTServer {
};
let hash_hex = hex::encode(info_hash_arr);
// 使用分片布隆过滤器进行高效去重
let is_duplicate = self.seen_hashes.check_and_set(&info_hash_arr);
if is_duplicate {
let dup_cb = self.on_duplicate.read().unwrap().clone();
if let Some(cb) = dup_cb {
let hash_hex_clone = hash_hex.clone();
tokio::spawn(async move {
cb(&hash_hex_clone);
});
}
return Ok(());
}
let filter_cb = self.filter.read().unwrap().clone();
if let Some(f) = filter_cb {
if !f(&hash_hex) { return Ok(()); }
src/sharded.rs
-4
View File
@@ -1,8 +1,6 @@
// 分片锁实现 - 大幅减少锁竞争，提升并发性能
//
// 核心思想：1个大锁 → N个小锁
// 性能提升：预期 3-4 倍
use bloomfilter::Bloom;
use std::collections::{HashSet, VecDeque};
use std::net::SocketAddr;
@@ -226,8 +224,6 @@ impl ShardedNodeQueue {
}
/// 获取随机节点（用于DHT响应）
/// 🚀 优化：IPv4 和 IPv6 分开存储，直接从对应队列获取，无需过滤
///
/// # Arguments
/// * `count` - 需要获取的节点数量
/// * `filter_ipv6` - 如果为 `Some(true)`，只返回 IPv6 节点；如果为 `Some(false)`，只返回 IPv4 节点；如果为 `None`，返回所有节点（混合）
src/types.rs
+9 -4
View File
@@ -82,21 +82,26 @@ pub struct DHTOptions {
/// 网络模式配置（仅IPv4、仅IPv6、或双栈）
pub netmode: NetMode,
/// 节点队列容量（默认 100000）
pub node_queue_capacity: usize,
}
impl Default for DHTOptions {
fn default() -> Self {
Self {
port: 0,
port: 6881, // BitTorrent DHT 默认端口
auto_metadata: true,
// 缩短超时，快速失败，不等待慢节点
metadata_timeout: 10,
metadata_timeout: 3,
// 加大队列，防止流量高峰丢包
max_metadata_queue_size: 10000,
max_metadata_queue_size: 100000,
// 提高并发，模拟 Node.js 的高并发 IO
max_metadata_worker_count: 1000,
// 默认双栈
netmode: NetMode::DualStack,
netmode: NetMode::Ipv4Only,
// 节点队列容量：100000 个节点（扩容以适应 DHT 网络裂变速度）
node_queue_capacity: 100000,
}
}
}