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

init

Browse Source
This commit is contained in:
桥下红药
2026-01-16 01:17:03 +08:00
parent 03ebbcc5e0
commit 7453571ecf
6 changed files with 468 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+2 -2
View File
@@ -8,7 +8,7 @@ pub mod scheduler; // 元数据调度器
pub use error::{DHTError, Result};
pub use server::{DHTServer, HashDiscovered};
pub use types::{DHTOptions, FileInfo, TorrentInfo};
pub use types::{DHTOptions, FileInfo, TorrentInfo, NetMode};
pub use sharded::{ShardedBloom, ShardedNodeQueue, NodeTuple};
pub use scheduler::MetadataScheduler;
@@ -16,6 +16,6 @@ pub use scheduler::MetadataScheduler;
pub mod prelude {
pub use crate::error::{DHTError, Result};
pub use crate::server::DHTServer;
pub use crate::types::{DHTOptions, FileInfo, TorrentInfo};
pub use crate::types::{DHTOptions, FileInfo, TorrentInfo, NetMode};
pub use crate::scheduler::MetadataScheduler;
}
src/main.rs
+2 -1
View File
@@ -21,11 +21,12 @@ async fn main() -> Result<()> {
.init();
let options = DHTOptions {
port: 45452,
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（双栈，默认）
};
// 统计计数器
src/protocol.rs
+2
View File
@@ -29,5 +29,7 @@ pub struct DhtResponse {
pub id: Option<serde_bytes::ByteBuf>,
#[serde(default)]
pub nodes: Option<serde_bytes::ByteBuf>,
#[serde(default)]
pub nodes6: Option<serde_bytes::ByteBuf>,
}
src/server.rs
+249 -37
View File
@@ -2,12 +2,12 @@ use crate::error::Result;
use crate::metadata::RbitFetcher;
use crate::protocol::{DhtMessage, DhtArgs, DhtResponse};
use crate::scheduler::MetadataScheduler;
use crate::types::{DHTOptions, TorrentInfo};
use crate::types::{DHTOptions, TorrentInfo, NetMode};
use crate::sharded::{ShardedBloom, ShardedNodeQueue, NodeTuple};
use rand::Rng;
use ahash::AHasher;
use std::hash::{Hash, Hasher};
use std::net::{IpAddr, SocketAddr};
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
@@ -49,6 +49,7 @@ pub struct DHTServer {
options: DHTOptions,
node_id: Vec<u8>,
socket: Arc<UdpSocket>,
socket_v6: Option<Arc<UdpSocket>>,
token_secret: Vec<u8>,
// 这些回调现在与 MetadataScheduler 共享
@@ -71,19 +72,69 @@ pub struct DHTServer {
impl DHTServer {
pub async fn new(options: DHTOptions) -> Result<Self> {
let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
#[cfg(not(windows))]
{ let _ = socket.set_reuse_port(true); }
let _ = socket.set_reuse_address(true);
socket.set_nonblocking(true)?;
// 增加网络缓冲区以应对高QPS
let _ = socket.set_recv_buffer_size(32 * 1024 * 1024); // 32MB（原16MB）
let _ = socket.set_send_buffer_size(8 * 1024 * 1024); // 8MB（原4MB）
let (socket, socket_v6) = match options.netmode {
NetMode::Ipv4Only => {
let sock = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
#[cfg(not(windows))]
{ let _ = sock.set_reuse_port(true); }
let _ = sock.set_reuse_address(true);
sock.set_nonblocking(true)?;
// 增加网络缓冲区以应对高QPS
let _ = sock.set_recv_buffer_size(32 * 1024 * 1024); // 32MB
let _ = sock.set_send_buffer_size(8 * 1024 * 1024); // 8MB
let addr: SocketAddr = format!("0.0.0.0:{}", options.port).parse().unwrap();
socket.bind(&addr.into())?;
let socket = UdpSocket::from_std(socket.into())?;
let addr: SocketAddr = format!("0.0.0.0:{}", options.port).parse().unwrap();
sock.bind(&addr.into())?;
(Arc::new(UdpSocket::from_std(sock.into())?), None)
},
NetMode::Ipv6Only => {
let sock = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))?;
#[cfg(not(windows))]
{ let _ = sock.set_reuse_port(true); }
let _ = sock.set_reuse_address(true);
// 设置仅IPv6模式（Windows默认是仅IPv6，Linux/Unix需要设置）
#[cfg(not(windows))]
{ let _ = sock.set_only_v6(true); }
sock.set_nonblocking(true)?;
let _ = sock.set_recv_buffer_size(32 * 1024 * 1024);
let _ = sock.set_send_buffer_size(8 * 1024 * 1024);
let addr: SocketAddr = format!("[::]:{}", options.port).parse().unwrap();
sock.bind(&addr.into())?;
(Arc::new(UdpSocket::from_std(sock.into())?), None)
},
NetMode::DualStack => {
// IPv4 socket
let sock_v4 = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
#[cfg(not(windows))]
{ let _ = sock_v4.set_reuse_port(true); }
let _ = sock_v4.set_reuse_address(true);
sock_v4.set_nonblocking(true)?;
let _ = sock_v4.set_recv_buffer_size(32 * 1024 * 1024);
let _ = sock_v4.set_send_buffer_size(8 * 1024 * 1024);
let addr_v4: SocketAddr = format!("0.0.0.0:{}", options.port).parse().unwrap();
sock_v4.bind(&addr_v4.into())?;
let socket = Arc::new(UdpSocket::from_std(sock_v4.into())?);
// IPv6 socket
let sock_v6 = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))?;
#[cfg(not(windows))]
{ let _ = sock_v6.set_reuse_port(true); }
let _ = sock_v6.set_reuse_address(true);
#[cfg(not(windows))]
{ let _ = sock_v6.set_only_v6(true); } // 仅IPv6，避免与IPv4冲突
sock_v6.set_nonblocking(true)?;
let _ = sock_v6.set_recv_buffer_size(32 * 1024 * 1024);
let _ = sock_v6.set_send_buffer_size(8 * 1024 * 1024);
let addr_v6: SocketAddr = format!("[::]:{}", options.port).parse().unwrap();
sock_v6.bind(&addr_v6.into())?;
let socket_v6 = Some(Arc::new(UdpSocket::from_std(sock_v6.into())?));
(socket, socket_v6)
},
};
let node_id = generate_random_id();
let mut rng = rand::thread_rng();
@@ -129,7 +180,8 @@ impl DHTServer {
let server = Self {
options: options.clone(),
node_id: node_id.clone(),
socket: Arc::new(socket),
socket,
socket_v6,
token_secret,
callback,
on_metadata_fetch,
@@ -149,6 +201,41 @@ impl DHTServer {
Ok(self.socket.local_addr()?)
}
/// 验证地址类型是否与当前 netmode 配置匹配
///
/// 防御性编程：虽然 socket 层面理论上不应该接收到不匹配的数据包，
/// 但在某些特殊情况下（如系统配置、双栈模式切换等）可能会有问题。
/// 此方法确保在应用层也进行验证，避免处理不匹配的地址类型。
fn is_addr_allowed(&self, addr: &SocketAddr) -> bool {
match self.options.netmode {
NetMode::Ipv4Only => addr.is_ipv4(),
NetMode::Ipv6Only => addr.is_ipv6(),
NetMode::DualStack => true, // 双栈模式接受所有地址类型
}
}
/// 根据目标地址选择合适的socket
fn select_socket(&self, addr: &SocketAddr) -> &Arc<UdpSocket> {
match self.options.netmode {
NetMode::Ipv4Only => {
// IPv4Only 模式：只有 IPv4 socket
&self.socket
},
NetMode::Ipv6Only => {
// IPv6Only 模式：只有 IPv6 socket
&self.socket
},
NetMode::DualStack => {
// 双栈模式：根据地址类型选择
if addr.is_ipv6() {
self.socket_v6.as_ref().unwrap_or(&self.socket)
} else {
&self.socket
}
},
}
}
/// 设置元数据获取前的检查回调
///
/// 此回调在发现新的 info_hash 后，但在实际连接对等端获取元数据之前执行。
@@ -160,7 +247,7 @@ impl DHTServer {
/// - 如果回调执行过慢，可能会导致任务队列堆积。
///
/// # 示例
/// ```rust,ignore
/// ```rust
/// server.on_metadata_fetch(|hash| async move {
/// // 检查数据库是否存在
/// // let exists = db.has(hash).await;
@@ -188,7 +275,7 @@ impl DHTServer {
/// - 否则会阻塞当前的元数据获取 Worker，降低系统吞吐量。
///
/// # 示例
/// ```rust,ignore
/// ```rust
/// server.on_torrent(|info| {
/// // 简单操作可以直接做
/// println!("Got torrent: {}", info.name);
@@ -250,17 +337,14 @@ impl DHTServer {
let mut loop_tick = 0;
loop {
// 自适应爬取速度：根据 Metadata 队列负载调整爬取策略
// 🚀 自适应爬取速度：根据 Metadata 队列负载调整爬取策略
let queue_len = server.metadata_queue_len.load(Ordering::Relaxed);
let queue_pressure = queue_len as f64 / server.max_metadata_queue_size as f64;
// 动态计算批次大小和休眠时间
let (batch_size, sleep_duration) = if queue_pressure < 0.5 {
// 🟢 绿区：队列空闲，全速爬取
(200, Duration::from_millis(10))
} else if queue_pressure < 0.8 {
let (batch_size, sleep_duration) = if queue_pressure < 0.8 {
// 🟡 黄区：队列有压力，适度减速
(200, Duration::from_millis(20))
(200, Duration::from_millis(10))
} else if queue_pressure < 0.95 {
// 🟠 橙区：队列高压，大幅减速
(20, Duration::from_millis(500))
@@ -269,11 +353,21 @@ impl DHTServer {
(0, Duration::from_millis(1000))
};
// 根据配置决定从哪个队列获取节点
let filter_ipv6 = match server.options.netmode {
NetMode::Ipv4Only => Some(false),
NetMode::Ipv6Only => Some(true),
NetMode::DualStack => None,
};
// 检查对应队列是否为空
let queue_empty = server.node_queue.is_empty_for(filter_ipv6);
let nodes_batch = {
if server.node_queue.is_empty() || batch_size == 0 {
if queue_empty || batch_size == 0 {
None
} else {
Some(server.node_queue.pop_batch(batch_size))
Some(server.node_queue.pop_batch(batch_size, filter_ipv6))
}
};
@@ -309,6 +403,7 @@ impl DHTServer {
fn start_receiver(&self) {
let socket = self.socket.clone();
let socket_v6 = self.socket_v6.clone();
let server = self.clone();
let num_workers = std::thread::available_parallelism()
@@ -331,6 +426,7 @@ impl DHTServer {
});
}
let senders_for_v6 = senders.clone();
tokio::spawn(async move {
let mut buf = [0u8; 65536];
let mut next_worker_idx = 0;
@@ -340,7 +436,6 @@ impl DHTServer {
Ok((size, addr)) => {
// 🛡️ 安全检查1：拒绝异常大的包（DHT 消息通常 < 2KB）
if size > 8192 {
#[cfg(debug_assertions)]
log::trace!("⚠️ 拒绝异常大的 UDP 包: {} 字节 from {}", size, addr);
continue;
}
@@ -371,9 +466,52 @@ impl DHTServer {
}
}
});
// IPv6 接收任务
if let Some(socket_v6) = socket_v6 {
let senders_v6 = senders_for_v6;
tokio::spawn(async move {
let mut buf = [0u8; 65536];
let mut next_worker_idx = 0;
loop {
match socket_v6.recv_from(&mut buf).await {
Ok((size, addr)) => {
if size > 8192 { continue; }
if size == 0 || buf[0] != b'd' { continue; }
let data = buf[..size].to_vec();
let tx = &senders_v6[next_worker_idx];
next_worker_idx = (next_worker_idx + 1) % num_workers;
match tx.try_send((data, addr)) {
Ok(_) => {},
Err(mpsc::error::TrySendError::Full(_)) => {
#[cfg(debug_assertions)]
log::trace!("UDP worker queue full, dropping packet");
},
Err(_) => { break; }
}
}
Err(_e) => {
tokio::time::sleep(Duration::from_millis(1)).await;
}
}
}
});
}
}
async fn handle_message(&self, data: &[u8], addr: SocketAddr) -> Result<()> {
// 🛡️ 验证地址类型是否与当前 netmode 配置匹配
// 防御性编程：虽然 socket 层面理论上不应该接收到不匹配的数据包，
// 但在某些特殊情况下（如系统配置、双栈模式切换等）可能会有问题
if !self.is_addr_allowed(&addr) {
log::trace!("⚠️ 拒绝不匹配的地址类型: {} (当前模式: {:?})", addr, self.options.netmode);
return Ok(());
}
let msg: DhtMessage = match serde_bencode::from_bytes(data) {
Ok(m) => m,
Err(_) => return Ok(()),
@@ -469,7 +607,7 @@ impl DHTServer {
// 使用 try_send，队列满时直接丢弃（背压）
if let Err(_) = self.hash_tx.try_send(event) {
#[cfg(debug_assertions)]
log::trace!("⚠️ Hash 队列满，丢弃 hash");
log::debug!("⚠️ Hash 队列满，丢弃 hash");
}
}
}
@@ -477,13 +615,23 @@ impl DHTServer {
}
async fn handle_response(&self, response: &DhtResponse) -> Result<()> {
// 处理 IPv4 节点
if let Some(nodes_bytes) = &response.nodes {
self.process_compact_nodes(nodes_bytes);
}
// 处理 IPv6 节点
if let Some(nodes6_bytes) = &response.nodes6 {
self.process_compact_nodes_v6(nodes6_bytes);
}
Ok(())
}
fn process_compact_nodes(&self, nodes_bytes: &[u8]) {
// 根据配置决定是否处理IPv4节点
if self.options.netmode == NetMode::Ipv6Only {
return;
}
if nodes_bytes.len() % 26 != 0 { return; }
// 使用分片队列，直接并发插入（无锁竞争）
@@ -498,6 +646,29 @@ impl DHTServer {
}
}
fn process_compact_nodes_v6(&self, nodes_bytes: &[u8]) {
// 根据配置决定是否处理IPv6节点
if self.options.netmode == NetMode::Ipv4Only {
return;
}
if nodes_bytes.len() % 38 != 0 { return; }
for chunk in nodes_bytes.chunks(38) {
let id = chunk[0..20].to_vec();
let port = u16::from_be_bytes([chunk[36], chunk[37]]);
let ip_bytes: [u8; 16] = match chunk[20..36].try_into() {
Ok(b) => b,
Err(_) => continue, // 如果转换失败（理论上不会），跳过该节点
};
let ip = Ipv6Addr::from(ip_bytes);
// 过滤掉不可用地址 (组播, 未指定等)
if !ip.is_unspecified() && !ip.is_multicast() {
let addr = SocketAddr::new(IpAddr::V6(ip), port);
self.node_queue.push(NodeTuple { id, addr });
}
}
}
async fn send_response(
&self,
tid: &[u8],
@@ -520,20 +691,50 @@ impl DHTServer {
r_dict.insert(b"token".to_vec(), serde_bencode::value::Value::Bytes(token));
if query_type == "get_peers" || query_type == "find_node" {
// 使用分片队列获取随机节点（无锁竞争）
let nodes = self.node_queue.get_random_nodes(8);
// 根据配置和请求方IP类型决定需要获取的节点类型
let requestor_is_ipv6 = addr.is_ipv6();
let filter_ipv6 = match self.options.netmode {
NetMode::Ipv4Only => Some(false), // 只要 IPv4
NetMode::Ipv6Only => Some(true), // 只要 IPv6
NetMode::DualStack => Some(requestor_is_ipv6), // 双栈模式：根据请求方IP类型返回对应类型的节点
};
// 使用分片队列获取随机节点（无锁竞争，带地址族过滤）
let nodes = self.node_queue.get_random_nodes(8, filter_ipv6);
let mut nodes_data = Vec::new();
let mut nodes6_data = Vec::new();
for node in nodes {
nodes_data.extend_from_slice(&node.id);
match node.addr.ip() {
IpAddr::V4(ip) => nodes_data.extend_from_slice(&ip.octets()),
_ => continue,
// IPv4 节点
IpAddr::V4(ip) => {
nodes_data.extend_from_slice(&node.id);
nodes_data.extend_from_slice(&ip.octets());
nodes_data.extend_from_slice(&node.addr.port().to_be_bytes());
},
// IPv6 节点
IpAddr::V6(ip) => {
nodes6_data.extend_from_slice(&node.id);
nodes6_data.extend_from_slice(&ip.octets());
nodes6_data.extend_from_slice(&node.addr.port().to_be_bytes());
},
}
nodes_data.extend_from_slice(&node.addr.port().to_be_bytes());
}
r_dict.insert(b"nodes".to_vec(), serde_bencode::value::Value::Bytes(nodes_data));
// 根据请求方IP类型返回对应类型的节点
// 在单栈模式下，get_random_nodes 已经过滤了节点类型，所以这里直接根据请求方类型返回即可
if requestor_is_ipv6 {
// 请求方是IPv6：返回IPv6节点
if !nodes6_data.is_empty() {
r_dict.insert(b"nodes6".to_vec(), serde_bencode::value::Value::Bytes(nodes6_data));
}
} else {
// 请求方是IPv4：返回IPv4节点
if !nodes_data.is_empty() {
r_dict.insert(b"nodes".to_vec(), serde_bencode::value::Value::Bytes(nodes_data));
}
}
}
let mut response: std::collections::HashMap<String, serde_bencode::value::Value> = std::collections::HashMap::new();
@@ -542,7 +743,7 @@ impl DHTServer {
response.insert("r".to_string(), serde_bencode::value::Value::Dict(r_dict));
if let Ok(encoded) = serde_bencode::to_bytes(&response) {
let _ = self.socket.send_to(&encoded, addr).await;
let _ = self.select_socket(&addr).send_to(&encoded, addr).await;
}
Ok(())
}
@@ -553,7 +754,18 @@ impl DHTServer {
match tokio::net::lookup_host(node).await {
Ok(addrs) => {
for addr in addrs {
if addr.is_ipv6() { continue; }
// 根据配置过滤地址
match self.options.netmode {
NetMode::Ipv4Only => {
if addr.is_ipv6() { continue; }
},
NetMode::Ipv6Only => {
if addr.is_ipv4() { continue; }
},
NetMode::DualStack => {
// 双栈模式，接受所有地址
},
}
let _ = self.send_find_node(addr, &target, &self.node_id).await;
}
}
@@ -574,7 +786,7 @@ impl DHTServer {
msg.insert("a".to_string(), serde_bencode::value::Value::Dict(args));
if let Ok(encoded) = serde_bencode::to_bytes(&msg) {
let _ = self.socket.send_to(&encoded, addr).await;
let _ = self.select_socket(&addr).send_to(&encoded, addr).await;
}
Ok(())
}
src/sharded.rs
+191 -86
View File
@@ -129,9 +129,10 @@ impl NodeQueueShard {
}
}
/// 分片节点队列 - 支持高并发
/// 分片节点队列 - 支持高并发，IPv4 和 IPv6 节点分开存储
pub struct ShardedNodeQueue {
shards: Vec<Mutex<NodeQueueShard>>,
shards_v4: Vec<Mutex<NodeQueueShard>>, // IPv4 节点分片
shards_v6: Vec<Mutex<NodeQueueShard>>, // IPv6 节点分片
}
impl ShardedNodeQueue {
@@ -139,134 +140,238 @@ impl ShardedNodeQueue {
pub fn new(total_capacity: usize) -> Self {
let capacity_per_shard = (total_capacity + QUEUE_SHARD_COUNT - 1) / QUEUE_SHARD_COUNT;
let shards = (0..QUEUE_SHARD_COUNT)
let shards_v4 = (0..QUEUE_SHARD_COUNT)
.map(|_| Mutex::new(NodeQueueShard::new(capacity_per_shard)))
.collect();
Self { shards }
let shards_v6 = (0..QUEUE_SHARD_COUNT)
.map(|_| Mutex::new(NodeQueueShard::new(capacity_per_shard)))
.collect();
Self { shards_v4, shards_v6 }
}
/// 添加节点
/// 添加节点（根据地址类型自动存入对应队列）
pub fn push(&self, node: NodeTuple) {
let shard_idx = self.addr_to_shard(&node.addr);
let mut shard = self.shards[shard_idx].lock().unwrap();
shard.push(node);
if node.addr.is_ipv6() {
let mut shard = self.shards_v6[shard_idx].lock().unwrap();
shard.push(node);
} else {
let mut shard = self.shards_v4[shard_idx].lock().unwrap();
shard.push(node);
}
}
/// 批量弹出节点
pub fn pop_batch(&self, count: usize) -> Vec<NodeTuple> {
///
/// # Arguments
/// * `count` - 需要获取的节点数量
/// * `filter_ipv6` - 如果为 `Some(true)`，只从 IPv6 队列获取；如果为 `Some(false)`，只从 IPv4 队列获取；如果为 `None`，从两个队列混合获取
pub fn pop_batch(&self, count: usize, filter_ipv6: Option<bool>) -> Vec<NodeTuple> {
let mut result = Vec::with_capacity(count);
let per_shard = (count + QUEUE_SHARD_COUNT - 1) / QUEUE_SHARD_COUNT;
// 从所有分片获取
for shard in &self.shards {
if result.len() >= count {
break;
}
let mut s = shard.lock().unwrap();
let nodes = s.pop_batch(per_shard);
result.extend(nodes);
match filter_ipv6 {
Some(true) => {
// 只从 IPv6 队列获取
for shard in &self.shards_v6 {
if result.len() >= count {
break;
}
let mut s = shard.lock().unwrap();
let nodes = s.pop_batch(per_shard);
result.extend(nodes);
}
},
Some(false) => {
// 只从 IPv4 队列获取
for shard in &self.shards_v4 {
if result.len() >= count {
break;
}
let mut s = shard.lock().unwrap();
let nodes = s.pop_batch(per_shard);
result.extend(nodes);
}
},
None => {
// 混合模式：从两个队列交替获取
for i in 0..QUEUE_SHARD_COUNT {
if result.len() >= count {
break;
}
// 从 IPv4 分片获取
let mut s4 = self.shards_v4[i].lock().unwrap();
let nodes4 = s4.pop_batch(per_shard / 2);
result.extend(nodes4);
drop(s4);
if result.len() >= count {
break;
}
// 从 IPv6 分片获取
let mut s6 = self.shards_v6[i].lock().unwrap();
let nodes6 = s6.pop_batch(per_shard / 2);
result.extend(nodes6);
drop(s6);
}
},
}
result
}
/// 获取随机节点（用于DHT响应）
/// 🚀 优化：使用储层采样算法，O(n)时间，无需clone全部节点
pub fn get_random_nodes(&self, count: usize) -> Vec<NodeTuple> {
/// 🚀 优化：IPv4 和 IPv6 分开存储，直接从对应队列获取，无需过滤
///
/// # Arguments
/// * `count` - 需要获取的节点数量
/// * `filter_ipv6` - 如果为 `Some(true)`，只返回 IPv6 节点；如果为 `Some(false)`，只返回 IPv4 节点；如果为 `None`，返回所有节点（混合）
pub fn get_random_nodes(&self, count: usize, filter_ipv6: Option<bool>) -> Vec<NodeTuple> {
match filter_ipv6 {
Some(true) => {
// 只要 IPv6 节点
self.get_random_nodes_from_shards(&self.shards_v6, count)
},
Some(false) => {
// 只要 IPv4 节点
self.get_random_nodes_from_shards(&self.shards_v4, count)
},
None => {
// 混合模式：从两个队列各取一半
let count_v4 = count / 2;
let count_v6 = count - count_v4;
let mut result = Vec::with_capacity(count);
result.extend(self.get_random_nodes_from_shards(&self.shards_v4, count_v4));
result.extend(self.get_random_nodes_from_shards(&self.shards_v6, count_v6));
result
},
}
}
/// 从指定的分片组中获取随机节点
fn get_random_nodes_from_shards(&self, shards: &[Mutex<NodeQueueShard>], count: usize) -> Vec<NodeTuple> {
use rand::Rng;
let mut rng = rand::thread_rng();
// 🚀 策略1：小规模请求用快速路径（最常见：8个节点）
if count <= 16 {
return self.get_random_nodes_fast(count);
}
// 🚀 策略2：大规模请求用储层采样
let mut result = Vec::with_capacity(count);
let mut seen = 0usize;
// 储层采样算法
for shard in &self.shards {
let s = shard.lock().unwrap();
let mut result = Vec::with_capacity(count);
let per_shard = (count + QUEUE_SHARD_COUNT - 1) / QUEUE_SHARD_COUNT;
for node in s.queue.iter() {
seen += 1;
for shard in shards {
if result.len() >= count {
break;
}
if result.len() < count {
// 前 count 个直接加入
result.push(node.clone());
} else {
// 后续以 count/seen 的概率替换
let j = rng.gen_range(0..seen);
if j < count {
result[j] = node.clone();
let s = shard.lock().unwrap();
let shard_len = s.queue.len();
if shard_len == 0 {
continue;
}
// 从当前分片随机选择最多 per_shard 个节点
let to_take = per_shard.min(shard_len).min(count - result.len());
// 生成随机索引（不重复）
let mut indices: Vec<usize> = (0..shard_len).collect();
// 只 shuffle 前 to_take 个（部分 shuffle，Fisher-Yates 优化）
for i in 0..to_take {
let j = rng.gen_range(i..shard_len);
indices.swap(i, j);
}
// 取前 to_take 个索引对应的节点
for i in 0..to_take {
if let Some(node) = s.queue.get(indices[i]) {
result.push(node.clone());
}
}
}
}
result
}
/// 快速路径：小规模随机选择（针对常见的8节点请求）
fn get_random_nodes_fast(&self, count: usize) -> Vec<NodeTuple> {
use rand::Rng;
let mut rng = rand::thread_rng();
let mut result = Vec::with_capacity(count);
// 从每个分片随机选择几个节点
let per_shard = (count + QUEUE_SHARD_COUNT - 1) / QUEUE_SHARD_COUNT;
for shard in &self.shards {
if result.len() >= count {
break;
}
let s = shard.lock().unwrap();
let shard_len = s.queue.len();
result
} else {
// 🚀 策略2：大规模请求用储层采样
let mut result = Vec::with_capacity(count);
let mut seen = 0usize;
if shard_len == 0 {
continue;
}
// 从当前分片随机选择最多 per_shard 个节点
let to_take = per_shard.min(shard_len).min(count - result.len());
// 生成随机索引（不重复）
let mut indices: Vec<usize> = (0..shard_len).collect();
// 只 shuffle 前 to_take 个（部分 shuffle，Fisher-Yates 优化）
for i in 0..to_take {
let j = rng.gen_range(i..shard_len);
indices.swap(i, j);
}
// 取前 to_take 个索引对应的节点
for i in 0..to_take {
if let Some(node) = s.queue.get(indices[i]) {
result.push(node.clone());
// 储层采样算法
for shard in shards {
let s = shard.lock().unwrap();
for node in s.queue.iter() {
seen += 1;
if result.len() < count {
// 前 count 个直接加入
result.push(node.clone());
} else {
// 后续以 count/seen 的概率替换
let j = rng.gen_range(0..seen);
if j < count {
result[j] = node.clone();
}
}
}
}
result
}
result
}
/// 获取总长度
/// 获取总长度（IPv4 + IPv6）
pub fn len(&self) -> usize {
self.shards
let len_v4: usize = self.shards_v4
.iter()
.map(|shard| shard.lock().unwrap().len())
.sum()
.sum();
let len_v6: usize = self.shards_v6
.iter()
.map(|shard| shard.lock().unwrap().len())
.sum();
len_v4 + len_v6
}
/// 检查是否为空
pub fn is_empty(&self) -> bool {
self.shards
let empty_v4 = self.shards_v4
.iter()
.all(|shard| shard.lock().unwrap().is_empty())
.all(|shard| shard.lock().unwrap().is_empty());
let empty_v6 = self.shards_v6
.iter()
.all(|shard| shard.lock().unwrap().is_empty());
empty_v4 && empty_v6
}
/// 检查指定地址族的队列是否为空
///
/// # Arguments
/// * `filter_ipv6` - 如果为 `Some(true)`，检查 IPv6 队列；如果为 `Some(false)`，检查 IPv4 队列；如果为 `None`，检查两个队列
pub fn is_empty_for(&self, filter_ipv6: Option<bool>) -> bool {
match filter_ipv6 {
Some(true) => {
// 检查 IPv6 队列
self.shards_v6
.iter()
.all(|shard| shard.lock().unwrap().is_empty())
},
Some(false) => {
// 检查 IPv4 队列
self.shards_v4
.iter()
.all(|shard| shard.lock().unwrap().is_empty())
},
None => self.is_empty(),
}
}
/// 根据地址计算分片索引
src/types.rs
+22
View File
@@ -1,5 +1,22 @@
use serde::{Deserialize, Serialize};
/// 网络模式配置
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NetMode {
/// 仅使用 IPv4
Ipv4Only,
/// 仅使用 IPv6
Ipv6Only,
/// 双栈（同时支持 IPv4 和 IPv6）
DualStack,
}
impl Default for NetMode {
fn default() -> Self {
Self::DualStack
}
}
/// 完整的种子信息（包含元数据）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TorrentInfo {
@@ -62,6 +79,9 @@ pub struct DHTOptions {
/// 并发元数据获取工作线程数
pub max_metadata_worker_count: usize,
/// 网络模式配置（仅IPv4、仅IPv6、或双栈）
pub netmode: NetMode,
}
impl Default for DHTOptions {
@@ -75,6 +95,8 @@ impl Default for DHTOptions {
max_metadata_queue_size: 10000,
// 提高并发，模拟 Node.js 的高并发 IO
max_metadata_worker_count: 1000,
// 默认双栈
netmode: NetMode::DualStack,
}
}
}