313 lines
11 KiB
Rust
313 lines
11 KiB
Rust
// 负责连接采集存储索引和接口层并定义应用级启动顺序
|
|
|
|
use std::{
|
|
str::FromStr,
|
|
sync::Arc,
|
|
time::{Duration, SystemTime, UNIX_EPOCH},
|
|
};
|
|
|
|
use dht_crawler::DHTServer;
|
|
use dht_search::{
|
|
domain::InfoHash,
|
|
search::SearchEngine,
|
|
storage::{RocksTorrentRepository, TorrentRepository},
|
|
};
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use crate::{
|
|
api::{self, ApiState},
|
|
config::AppConfig,
|
|
crawler::pipeline::PersistencePipeline,
|
|
error::AppError,
|
|
shutdown,
|
|
};
|
|
|
|
pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> {
|
|
std::fs::create_dir_all(&config.data_dir)?;
|
|
let database_path = config.data_dir.join("rocksdb");
|
|
let repository = Arc::new(RocksTorrentRepository::open(&database_path)?);
|
|
let (search, search_created) = SearchEngine::open_with_status(config.data_dir.join("tantivy"))?;
|
|
if search_created {
|
|
let records = repository.prepare_full_reindex()?;
|
|
tracing::info!(records, "检测到新搜索索引并准备全量重建");
|
|
}
|
|
let repository_api: Arc<dyn TorrentRepository> = repository.clone();
|
|
let mut persistence =
|
|
PersistencePipeline::start(repository_api, config.persistence_queue_capacity);
|
|
let ingress = persistence.ingress.clone();
|
|
|
|
let options = config.dht_options();
|
|
let server = DHTServer::new(options.clone()).await?;
|
|
server.on_error(|error| tracing::error!(%error, "DHT 运行时错误"));
|
|
|
|
let sampled_repository = repository.clone();
|
|
server.on_sampled_hashes(move |hashes| {
|
|
let repository = sampled_repository.clone();
|
|
async move {
|
|
let fallback = hashes.clone();
|
|
let info_hashes: Vec<_> = hashes.into_iter().map(InfoHash::from_bytes).collect();
|
|
let result = tokio::task::spawn_blocking(move || {
|
|
repository.filter_unknown_and_observe(&info_hashes, unix_timestamp())
|
|
})
|
|
.await;
|
|
match result {
|
|
Ok(Ok(unknown)) => unknown
|
|
.into_iter()
|
|
.map(|info_hash| *info_hash.as_bytes())
|
|
.collect(),
|
|
Ok(Err(error)) => {
|
|
tracing::error!(%error, "采样 infohash 批量持久化去重失败");
|
|
fallback
|
|
}
|
|
Err(error) => {
|
|
tracing::error!(%error, "采样 infohash 批量去重任务异常");
|
|
fallback
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
let gate_repository = repository.clone();
|
|
server.on_metadata_fetch(move |hash| {
|
|
let repository = gate_repository.clone();
|
|
async move {
|
|
let Ok(info_hash) = InfoHash::from_str(&hash) else {
|
|
tracing::warn!(%hash, "DHT 提供了无效 infohash");
|
|
return false;
|
|
};
|
|
let result = tokio::task::spawn_blocking(move || {
|
|
repository.observe_existing(info_hash, unix_timestamp())
|
|
})
|
|
.await;
|
|
match result {
|
|
Ok(Ok(already_exists)) => !already_exists,
|
|
Ok(Err(error)) => {
|
|
tracing::error!(%error, %hash, "持久化去重查询失败");
|
|
true
|
|
}
|
|
Err(error) => {
|
|
tracing::error!(%error, %hash, "持久化去重任务异常");
|
|
true
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
let callback_ingress = ingress.clone();
|
|
server.on_torrent_with_ack(move |torrent| callback_ingress.try_enqueue(torrent));
|
|
server.on_metadata_fetch_complete(|completion| {
|
|
tracing::debug!(
|
|
info_hash = %completion.info_hash,
|
|
status = ?completion.status,
|
|
attempts = completion.attempts,
|
|
"Metadata 任务完成"
|
|
)
|
|
});
|
|
|
|
tracing::info!(
|
|
dht_port = options.port,
|
|
data_dir = %config.data_dir.display(),
|
|
persistence_queue_capacity = config.persistence_queue_capacity,
|
|
"dht-search 启动"
|
|
);
|
|
|
|
let monitor_cancel = CancellationToken::new();
|
|
let monitor = tokio::spawn(monitor(
|
|
server.clone(),
|
|
ingress,
|
|
config.stats_interval_secs,
|
|
monitor_cancel.clone(),
|
|
));
|
|
let index_cancel = CancellationToken::new();
|
|
let (index_fatal_tx, mut index_fatal) = tokio::sync::oneshot::channel();
|
|
let index_task = tokio::spawn(run_indexer(
|
|
repository.clone(),
|
|
search.clone(),
|
|
config.index_batch_size,
|
|
Duration::from_millis(config.index_interval_millis),
|
|
index_cancel.clone(),
|
|
index_fatal_tx,
|
|
));
|
|
let api_cancel = CancellationToken::new();
|
|
let mut api_task = tokio::spawn(api::serve(
|
|
config.http.listen,
|
|
ApiState {
|
|
repository: repository.clone(),
|
|
search,
|
|
dht_stats: server.runtime_stats(),
|
|
persistence: persistence.ingress.clone(),
|
|
},
|
|
api_cancel.clone(),
|
|
));
|
|
|
|
let run_duration = async {
|
|
match config.run_duration_secs {
|
|
Some(seconds) => tokio::time::sleep(Duration::from_secs(seconds)).await,
|
|
None => std::future::pending().await,
|
|
}
|
|
};
|
|
tokio::pin!(run_duration);
|
|
|
|
let run_result = tokio::select! {
|
|
result = server.start() => result.map_err(AppError::from),
|
|
_ = shutdown::signal() => {
|
|
tracing::info!("收到退出信号");
|
|
Ok(())
|
|
}
|
|
_ = &mut run_duration => {
|
|
tracing::info!("达到配置的运行时长");
|
|
Ok(())
|
|
}
|
|
fatal = &mut persistence.fatal => {
|
|
let message = fatal.unwrap_or_else(|_| "持久化 worker 意外停止".to_owned());
|
|
Err(AppError::PersistenceWorker(message))
|
|
}
|
|
fatal = &mut index_fatal => {
|
|
let message = fatal.unwrap_or_else(|_| "索引 worker 意外停止".to_owned());
|
|
Err(AppError::IndexWorker(message))
|
|
}
|
|
result = &mut api_task => {
|
|
match result {
|
|
Ok(Ok(())) => Err(AppError::Config("HTTP 服务意外停止".to_owned())),
|
|
Ok(Err(error)) => Err(AppError::Io(error)),
|
|
Err(error) => Err(AppError::Config(format!("HTTP 服务任务异常: {error}"))),
|
|
}
|
|
}
|
|
};
|
|
|
|
server.shutdown();
|
|
monitor_cancel.cancel();
|
|
let _ = monitor.await;
|
|
persistence.close_and_join().await?;
|
|
index_cancel.cancel();
|
|
match index_task.await {
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(error)) => return Err(AppError::IndexWorker(error)),
|
|
Err(error) => return Err(AppError::IndexWorker(error.to_string())),
|
|
}
|
|
api_cancel.cancel();
|
|
if !api_task.is_finished() {
|
|
match api_task.await {
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(error)) => return Err(AppError::Io(error)),
|
|
Err(error) => return Err(AppError::Config(format!("HTTP 服务任务异常: {error}"))),
|
|
}
|
|
}
|
|
tracing::info!("dht-search 已安全停止");
|
|
run_result
|
|
}
|
|
|
|
async fn run_indexer(
|
|
repository: Arc<RocksTorrentRepository>,
|
|
search: SearchEngine,
|
|
batch_size: usize,
|
|
interval: Duration,
|
|
cancel: CancellationToken,
|
|
fatal: tokio::sync::oneshot::Sender<String>,
|
|
) -> Result<(), String> {
|
|
let mut ticker = tokio::time::interval(interval);
|
|
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
loop {
|
|
tokio::select! {
|
|
_ = cancel.cancelled() => break,
|
|
_ = ticker.tick() => {
|
|
let indexed = index_one_batch(repository.clone(), search.clone(), batch_size).await;
|
|
match indexed {
|
|
Ok(count) if count > 0 => tracing::debug!(count, "搜索索引已提交"),
|
|
Ok(_) => {}
|
|
Err(error) => {
|
|
let _ = fatal.send(error.clone());
|
|
return Err(error);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
loop {
|
|
let count = index_one_batch(repository.clone(), search.clone(), batch_size).await?;
|
|
if count == 0 {
|
|
break;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn index_one_batch(
|
|
repository: Arc<RocksTorrentRepository>,
|
|
search: SearchEngine,
|
|
batch_size: usize,
|
|
) -> Result<usize, String> {
|
|
tokio::task::spawn_blocking(move || {
|
|
search
|
|
.index_pending(repository.as_ref(), batch_size, unix_timestamp())
|
|
.map_err(|error| error.to_string())
|
|
})
|
|
.await
|
|
.map_err(|error| error.to_string())?
|
|
}
|
|
|
|
async fn monitor(
|
|
server: DHTServer,
|
|
ingress: crate::crawler::pipeline::PersistenceIngress,
|
|
interval_secs: u64,
|
|
cancel: CancellationToken,
|
|
) {
|
|
let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
|
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
let mut previous_udp_tx = 0;
|
|
let mut previous_metadata_attempts = 0;
|
|
loop {
|
|
tokio::select! {
|
|
_ = cancel.cancelled() => break,
|
|
_ = interval.tick() => {
|
|
let dht = server.runtime_stats().snapshot();
|
|
let observability = server.runtime_stats().observability_snapshot();
|
|
let storage = ingress.snapshot();
|
|
let udp_tx_per_second = observability
|
|
.udp_tx_packets
|
|
.saturating_sub(previous_udp_tx)
|
|
/ interval_secs;
|
|
let metadata_connects_per_second = dht
|
|
.metadata_peer_attempts
|
|
.saturating_sub(previous_metadata_attempts)
|
|
/ interval_secs;
|
|
previous_udp_tx = observability.udp_tx_packets;
|
|
previous_metadata_attempts = dht.metadata_peer_attempts;
|
|
tracing::info!(
|
|
nodes = dht.node_pool_size,
|
|
udp_tx = observability.udp_tx_packets,
|
|
udp_tx_per_second,
|
|
find_node_queries = dht.queries_new + dht.queries_revisit + dht.queries_bootstrap,
|
|
peer_lookup_queries = dht.peer_lookup_queries,
|
|
peer_lookup_preferred_succeeded = dht.peer_lookup_preferred_succeeded,
|
|
peer_lookup_fallbacks = dht.peer_lookup_fallbacks,
|
|
sample_queries = dht.sample_infohashes_queries,
|
|
sampled_hashes = dht.sample_infohashes_hashes_discovered,
|
|
sampled_hashes_filtered = dht.sample_infohashes_hashes_filtered,
|
|
peers = dht.peer_lookup_peers_found,
|
|
metadata_connects_per_second,
|
|
metadata_in_flight = dht.metadata_in_flight,
|
|
metadata_ok = dht.metadata_peer_succeeded,
|
|
metadata_failed = dht.metadata_peer_failed,
|
|
persistence_accepted = storage.accepted,
|
|
persistence_inserted = storage.inserted,
|
|
persistence_updated = storage.updated,
|
|
persistence_rejected_full = storage.rejected_full,
|
|
persistence_invalid = storage.invalid,
|
|
persistence_failed = storage.failed,
|
|
persistence_queue = storage.queue_depth,
|
|
"运行状态"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn unix_timestamp() -> u64 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs()
|
|
}
|