From 104b165f465e1b0d73de536f1435f7fcb6879a7a Mon Sep 17 00:00:00 2001 From: chuan Date: Mon, 10 Aug 2026 13:02:04 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E7=A3=81=E7=9B=98?= =?UTF-8?q?=E7=A9=BA=E9=97=B4=E5=8F=AA=E8=AF=BB=E4=BF=9D=E6=8A=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.lock | 11 + TODOS.md | 6 +- dht-search.example.toml | 6 + dht-search/Cargo.toml | 1 + dht-search/README.md | 15 ++ dht-search/src/api/handlers.rs | 10 + dht-search/src/api/mod.rs | 25 ++- dht-search/src/api/response.rs | 9 + dht-search/src/app.rs | 41 +++- dht-search/src/config.rs | 37 +++ dht-search/src/crawler/pipeline.rs | 49 +++- dht-search/src/disk_guard.rs | 346 +++++++++++++++++++++++++++++ dht-search/src/index_worker.rs | 42 +++- dht-search/src/main.rs | 1 + dht-search/src/monitor.rs | 9 +- dht-search/src/verification.rs | 44 +++- web/src/App.vue | 10 +- web/src/types/api.ts | 9 + 18 files changed, 641 insertions(+), 30 deletions(-) create mode 100644 dht-search/src/disk_guard.rs diff --git a/Cargo.lock b/Cargo.lock index 6733916..b179d89 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -656,6 +656,7 @@ dependencies = [ "blake3", "clap", "dht-crawler", + "fs2", "hex", "regex", "rmp-serde", @@ -817,6 +818,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "fs4" version = "0.13.1" diff --git a/TODOS.md b/TODOS.md index b96f687..855052a 100644 --- a/TODOS.md +++ b/TODOS.md @@ -282,7 +282,7 @@ - [ ] 记录 RocksDB block cache memtable 和 compaction 指标 - [ ] 记录 Tantivy IndexWriter 内存和 commit 延迟 - [ ] 根据实测调整批量大小队列容量和并发 -- [ ] 增加磁盘剩余空间保护和只读降级策略 +- [x] 增加带排空阶段恢复滞回和探测失败保护的磁盘只读降级策略 - [ ] 增加数据库备份检查点和恢复验证 - [ ] 增加日志轮转和保留策略 - [x] 验证间歇运行和正常退出恢复 @@ -294,7 +294,7 @@ - [ ] 内存使用在目标上限内稳定 - [ ] 队列和缓存不会随运行时间无限增长 -- [ ] 磁盘不足时能够安全停止写入 +- [x] 磁盘不足时能够拒绝新任务排空持久化队列并保留搜索能力 - [ ] 备份可以在独立目录恢复并搜索 - [ ] 连续运行期间没有数据格式损坏和不可恢复任务 @@ -336,4 +336,4 @@ 完成二十四小时持续运行并继续观察私有内存 Metadata 成功率候选队列深度和每条成功 Metadata 的网络成本 -随后增加磁盘剩余空间保护日志轮转和 RocksDB 检查点恢复 +随后增加日志轮转和 RocksDB 检查点恢复 diff --git a/dht-search.example.toml b/dht-search.example.toml index 40a6d06..4b4d55b 100644 --- a/dht-search.example.toml +++ b/dht-search.example.toml @@ -7,6 +7,12 @@ stats_interval_secs = 10 index_batch_size = 1024 index_interval_millis = 5000 +[disk_guard] +enabled = true +check_interval_secs = 10 +minimum_free_bytes = 5368709120 +resume_free_bytes = 6442450944 + [metadata_limits] max_metadata_bytes = 10485760 max_files = 20000 diff --git a/dht-search/Cargo.toml b/dht-search/Cargo.toml index 974d627..3bcbf40 100644 --- a/dht-search/Cargo.toml +++ b/dht-search/Cargo.toml @@ -19,6 +19,7 @@ blake3 = "1.8.5" clap = { version = "4.5", features = ["derive"] } dht-crawler = { path = "../dht-crawler", features = ["metrics"] } hex = "0.4" +fs2 = "0.4" rocksdb = { version = "0.24.0", default-features = false, features = ["bindgen-runtime", "lz4"], optional = true } rmp-serde = "1.3" regex = "1.12" diff --git a/dht-search/README.md b/dht-search/README.md index f4b4b36..f68e721 100644 --- a/dht-search/README.md +++ b/dht-search/README.md @@ -86,6 +86,21 @@ Metadata 下载和可用性握手共用 `metadata_connects_per_second` 预算不 Windows 下索引每五秒批量提交 临时文件占用会自动指数退避重试且不会停止采集 HTTP 服务或丢失 RocksDB 待索引状态 +### 磁盘空间保护 + +应用默认每十秒检查 `data_dir` 所在磁盘的剩余空间 低于保护阈值时先停止接收新 Metadata DHT 状态更新索引任务和可用性验证任务 已经进入有界持久化队列的记录会继续排空 随后进入只读保护 + +只读保护期间 RocksDB 和 Tantivy 不再产生业务写入 现有搜索详情健康检查和 Web 页面仍然可用 剩余空间达到独立恢复阈值后自动恢复采集 使用两个阈值可以避免临界空间附近反复暂停和恢复 + +| 配置项 | 默认值 | 作用 | +|---|---:|---| +| `disk_guard.enabled` | `true` | 是否启用磁盘空间保护 | +| `disk_guard.check_interval_secs` | `10` | 剩余空间检查间隔 | +| `disk_guard.minimum_free_bytes` | `5368709120` | 低于 5 GiB 时停止接收新任务 | +| `disk_guard.resume_free_bytes` | `6442450944` | 恢复到 6 GiB 时重新接受写入 | + +磁盘空间探测失败时采用保守策略进入保护状态 `/stats` 返回 `disk_state` `disk_available_bytes` 阈值 活跃写入数 探测失败数 状态转换数和拒绝任务数 Web 运行状态使用绿色或黄色状态点展示正常与保护状态 + ### Metadata 安全限制 应用会在 Metadata 下载和进入 RocksDB 前执行两层资源与结构校验 diff --git a/dht-search/src/api/handlers.rs b/dht-search/src/api/handlers.rs index 3a33b32..75d3a26 100644 --- a/dht-search/src/api/handlers.rs +++ b/dht-search/src/api/handlers.rs @@ -35,12 +35,22 @@ pub(crate) async fn stats(State(state): State) -> Json let dht = state.dht_stats.snapshot(); let observability = state.dht_stats.observability_snapshot(); let persistence = state.persistence.snapshot(); + let disk = state.disk_guard.snapshot(); let filtered = persistence.filtered; let verification = state .verification .as_ref() .map(|ingress| ingress.stats().snapshot()); Json(StatsResponse { + disk_state: disk.mode.as_str(), + disk_available_bytes: disk.available_bytes, + disk_minimum_free_bytes: disk.minimum_free_bytes, + disk_resume_free_bytes: disk.resume_free_bytes, + disk_active_writes: disk.active_writes, + disk_probe_failed: disk.probe_failed, + disk_probe_failures: disk.probe_failures, + disk_transitions: disk.transitions, + disk_rejected_new_work: disk.rejected_new_work, nodes: dht.node_pool_size, udp_tx_packets: observability.udp_tx_packets, find_node_queries: dht diff --git a/dht-search/src/api/mod.rs b/dht-search/src/api/mod.rs index f58cda4..52a6790 100644 --- a/dht-search/src/api/mod.rs +++ b/dht-search/src/api/mod.rs @@ -12,7 +12,9 @@ use dht_search::{search::SearchEngine, storage::TorrentRepository}; use tokio_util::sync::CancellationToken; use tower_http::services::{ServeDir, ServeFile}; -use crate::{crawler::pipeline::PersistenceIngress, verification::VerificationIngress}; +use crate::{ + crawler::pipeline::PersistenceIngress, disk_guard::DiskGuard, verification::VerificationIngress, +}; #[derive(Clone)] pub(crate) struct ApiState { @@ -21,6 +23,7 @@ pub(crate) struct ApiState { pub(crate) dht_stats: DhtRuntimeStats, pub(crate) persistence: PersistenceIngress, pub(crate) verification: Option, + pub(crate) disk_guard: DiskGuard, } pub(crate) async fn serve( @@ -67,7 +70,10 @@ mod tests { use tempfile::TempDir; use tower::ServiceExt; - use crate::{crawler::pipeline::PersistencePipeline, verification::VerificationIngress}; + use crate::{ + config::DiskGuardConfig, crawler::pipeline::PersistencePipeline, disk_guard::DiskGuard, + verification::VerificationIngress, + }; use super::*; @@ -105,8 +111,16 @@ mod tests { let search = SearchEngine::open(directory.path().join("tantivy")).unwrap(); search.index_pending(repository.as_ref(), 10, 20).unwrap(); let repository_trait: Arc = repository.clone(); - let persistence = - PersistencePipeline::start(repository_trait.clone(), 4, MetadataLimits::default()); + let disk_guard = DiskGuard::new(&DiskGuardConfig { + enabled: false, + ..DiskGuardConfig::default() + }); + let persistence = PersistencePipeline::start( + repository_trait.clone(), + 4, + MetadataLimits::default(), + disk_guard.clone(), + ); let verification = VerificationIngress::for_test(repository.clone(), 10); let web_dir = directory.path().join("web"); std::fs::create_dir_all(&web_dir).unwrap(); @@ -118,6 +132,7 @@ mod tests { dht_stats: DhtRuntimeStats::default(), persistence: persistence.ingress.clone(), verification: Some(verification), + disk_guard, }, web_dir, ); @@ -147,6 +162,8 @@ mod tests { .unwrap(); assert_eq!(json["metadata_filtered"], 0); assert_eq!(json["metadata_filtered_too_many_files"], 0); + assert_eq!(json["disk_state"], "normal"); + assert!(json["disk_available_bytes"].is_null()); let response = app .clone() diff --git a/dht-search/src/api/response.rs b/dht-search/src/api/response.rs index 17c9305..cea5b49 100644 --- a/dht-search/src/api/response.rs +++ b/dht-search/src/api/response.rs @@ -16,6 +16,15 @@ pub(crate) struct ErrorResponse { #[derive(Debug, Serialize)] pub(crate) struct StatsResponse { + pub(crate) disk_state: &'static str, + pub(crate) disk_available_bytes: Option, + pub(crate) disk_minimum_free_bytes: u64, + pub(crate) disk_resume_free_bytes: u64, + pub(crate) disk_active_writes: usize, + pub(crate) disk_probe_failed: bool, + pub(crate) disk_probe_failures: u64, + pub(crate) disk_transitions: u64, + pub(crate) disk_rejected_new_work: u64, pub(crate) nodes: usize, pub(crate) udp_tx_packets: u64, pub(crate) find_node_queries: u64, diff --git a/dht-search/src/app.rs b/dht-search/src/app.rs index c40a026..e8367e6 100644 --- a/dht-search/src/app.rs +++ b/dht-search/src/app.rs @@ -18,12 +18,14 @@ use crate::{ api::{self, ApiState}, config::AppConfig, crawler::pipeline::PersistencePipeline, + disk_guard::{self, DiskGuard}, error::AppError, index_worker, monitor, shutdown, verification, }; pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { std::fs::create_dir_all(&config.data_dir)?; + let disk_guard = DiskGuard::new(&config.disk_guard); let database_path = config.data_dir.join("rocksdb"); let metadata_limits = config.metadata_limits(); let content_filter = Arc::new(config.content_filter()?); @@ -38,29 +40,41 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { } else { SearchEngine::open_with_status(&search_path)? }; - if search_created { - let records = repository.prepare_full_reindex()?; - tracing::info!(records, "检测到新搜索索引并准备全量重建"); - } + disk_guard.probe(&config.data_dir, 0); let repository_api: Arc = repository.clone(); let mut persistence = PersistencePipeline::start( repository_api, config.persistence_queue_capacity, metadata_limits, + disk_guard.clone(), ); let ingress = persistence.ingress.clone(); + let disk_cancel = CancellationToken::new(); + let disk_task = tokio::spawn(disk_guard::run( + disk_guard.clone(), + config.data_dir.clone(), + config.disk_guard.clone(), + ingress.clone(), + disk_cancel.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(); + let sampled_disk_guard = disk_guard.clone(); server.on_sampled_hashes(move |hashes| { let repository = sampled_repository.clone(); + let disk_guard = sampled_disk_guard.clone(); async move { + let Some(permit) = disk_guard.begin_admission() else { + return Vec::new(); + }; let fallback = hashes.clone(); let info_hashes: Vec<_> = hashes.into_iter().map(InfoHash::from_bytes).collect(); let result = tokio::task::spawn_blocking(move || { + let _permit = permit; repository.filter_unknown_and_observe(&info_hashes, unix_timestamp()) }) .await; @@ -82,14 +96,20 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { }); let gate_repository = repository.clone(); + let gate_disk_guard = disk_guard.clone(); server.on_metadata_fetch(move |hash| { let repository = gate_repository.clone(); + let disk_guard = gate_disk_guard.clone(); async move { + let Some(permit) = disk_guard.begin_admission() else { + return false; + }; let Ok(info_hash) = InfoHash::from_str(&hash) else { tracing::warn!(%hash, "DHT 提供了无效 infohash"); return false; }; let result = tokio::task::spawn_blocking(move || { + let _permit = permit; repository.observe_existing(info_hash, unix_timestamp()) }) .await; @@ -126,6 +146,7 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { repository.clone(), server.clone(), config.verification.clone(), + disk_guard.clone(), verification_cancel.clone(), ); let cancel = verification_cancel.clone(); @@ -161,6 +182,7 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { let monitor = tokio::spawn(monitor::run( server.clone(), ingress, + disk_guard.clone(), config.stats_interval_secs, monitor_cancel.clone(), )); @@ -169,8 +191,12 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { let index_task = tokio::spawn(index_worker::run( repository.clone(), search.clone(), - config.index_batch_size, - Duration::from_millis(config.index_interval_millis), + index_worker::IndexWorkerOptions { + batch_size: config.index_batch_size, + interval: Duration::from_millis(config.index_interval_millis), + prepare_full_reindex: search_created, + }, + disk_guard.clone(), index_cancel.clone(), index_fatal_tx, )); @@ -184,6 +210,7 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { dht_stats: server.runtime_stats(), persistence: persistence.ingress.clone(), verification: verification_ingress, + disk_guard: disk_guard.clone(), }, api_cancel.clone(), )); @@ -237,6 +264,8 @@ pub(crate) async fn run(config: AppConfig) -> Result<(), AppError> { } } server.shutdown(); + disk_cancel.cancel(); + let _ = disk_task.await; monitor_cancel.cancel(); let _ = monitor.await; persistence.close_and_join().await?; diff --git a/dht-search/src/config.rs b/dht-search/src/config.rs index e31a578..4533ef5 100644 --- a/dht-search/src/config.rs +++ b/dht-search/src/config.rs @@ -35,6 +35,7 @@ pub(crate) struct AppConfig { pub(crate) index_interval_millis: u64, pub(crate) metadata_limits: MetadataLimitsConfig, pub(crate) dht: DhtConfig, + pub(crate) disk_guard: DiskGuardConfig, pub(crate) http: HttpConfig, pub(crate) verification: VerificationConfig, } @@ -70,6 +71,15 @@ pub(crate) struct HttpConfig { pub(crate) web_dir: PathBuf, } +#[derive(Debug, Clone, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub(crate) struct DiskGuardConfig { + pub(crate) enabled: bool, + pub(crate) check_interval_secs: u64, + pub(crate) minimum_free_bytes: u64, + pub(crate) resume_free_bytes: u64, +} + #[derive(Debug, Clone, Deserialize)] #[serde(default, deny_unknown_fields)] pub(crate) struct VerificationConfig { @@ -252,6 +262,14 @@ impl AppConfig { "验证队列容量并发尝试数租约和轮询间隔必须大于零".to_owned(), )); } + if self.disk_guard.check_interval_secs == 0 + || self.disk_guard.minimum_free_bytes == 0 + || self.disk_guard.resume_free_bytes <= self.disk_guard.minimum_free_bytes + { + return Err(AppError::Config( + "磁盘检查间隔必须大于零且恢复阈值必须大于保护阈值".to_owned(), + )); + } if self.dht.max_outbound_queries_per_second == 0 || self.dht.outbound_query_burst == 0 || self.dht.metadata_connects_per_second == 0 @@ -283,6 +301,7 @@ impl Default for AppConfig { index_interval_millis: 5_000, metadata_limits: MetadataLimitsConfig::default(), dht: DhtConfig::default(), + disk_guard: DiskGuardConfig::default(), http: HttpConfig::default(), verification: VerificationConfig::default(), } @@ -336,6 +355,17 @@ impl Default for HttpConfig { } } +impl Default for DiskGuardConfig { + fn default() -> Self { + Self { + enabled: true, + check_interval_secs: 10, + minimum_free_bytes: 5 * 1024 * 1024 * 1024, + resume_free_bytes: 6 * 1024 * 1024 * 1024, + } + } +} + impl Default for VerificationConfig { fn default() -> Self { Self { @@ -452,4 +482,11 @@ mod tests { config.dht.sample_new_node_percent = 101; assert!(matches!(config.validate(), Err(AppError::Config(_)))); } + + #[test] + fn disk_resume_threshold_must_exceed_minimum() { + let mut config = AppConfig::default(); + config.disk_guard.resume_free_bytes = config.disk_guard.minimum_free_bytes; + assert!(matches!(config.validate(), Err(AppError::Config(_)))); + } } diff --git a/dht-search/src/crawler/pipeline.rs b/dht-search/src/crawler/pipeline.rs index 4a6afc0..8d52572 100644 --- a/dht-search/src/crawler/pipeline.rs +++ b/dht-search/src/crawler/pipeline.rs @@ -17,6 +17,7 @@ use dht_search::{ }; use tokio::sync::oneshot; +use crate::disk_guard::DiskGuard; use crate::error::AppError; #[derive(Clone)] @@ -24,6 +25,7 @@ pub(crate) struct PersistenceIngress { sender: Arc>>>, stats: Arc, limits: MetadataLimits, + disk_guard: DiskGuard, } pub(crate) struct PersistencePipeline { @@ -71,15 +73,23 @@ impl PersistencePipeline { repository: Arc, capacity: usize, limits: MetadataLimits, + disk_guard: DiskGuard, ) -> Self { let (sender, receiver) = mpsc::sync_channel::(capacity); let (fatal_tx, fatal) = oneshot::channel(); let stats = Arc::new(PersistenceStats::default()); let worker_stats = stats.clone(); + let worker_disk_guard = disk_guard.clone(); let worker = thread::Builder::new() .name("torrent-persistence".to_owned()) .spawn(move || { while let Ok(item) = receiver.recv() { + let _permit = loop { + if let Some(permit) = worker_disk_guard.begin_drain_write() { + break permit; + } + thread::sleep(std::time::Duration::from_millis(250)); + }; worker_stats.queue_depth.fetch_sub(1, Ordering::Relaxed); let outcome = match item { PersistenceItem::Record(record) => repository.upsert(record).map(Some), @@ -111,6 +121,7 @@ impl PersistencePipeline { sender: Arc::new(Mutex::new(Some(sender))), stats, limits, + disk_guard, }, fatal, worker, @@ -129,6 +140,9 @@ impl PersistencePipeline { impl PersistenceIngress { pub(crate) fn try_enqueue(&self, torrent: TorrentInfo) -> bool { + let Some(_permit) = self.disk_guard.begin_admission() else { + return false; + }; let rejected_info_hash = InfoHash::from_str(&torrent.info_hash).ok(); let rejected_at = torrent.timestamp; let record = match TorrentRecord::try_from_with_limits(torrent, self.limits) { @@ -230,6 +244,14 @@ mod tests { }; use super::*; + use crate::config::DiskGuardConfig; + + fn disk_guard() -> DiskGuard { + DiskGuard::new(&DiskGuardConfig { + enabled: false, + ..DiskGuardConfig::default() + }) + } #[derive(Default)] struct MemoryRepository { @@ -348,7 +370,12 @@ mod tests { #[tokio::test] async fn accepted_record_is_drained_before_shutdown() { let repository = Arc::new(MemoryRepository::default()); - let pipeline = PersistencePipeline::start(repository.clone(), 1, MetadataLimits::default()); + let pipeline = PersistencePipeline::start( + repository.clone(), + 1, + MetadataLimits::default(), + disk_guard(), + ); assert!(pipeline.ingress.try_enqueue(torrent())); pipeline.close_and_join().await.unwrap(); let records = repository.records.lock().unwrap(); @@ -362,7 +389,7 @@ mod tests { max_files: 1, ..MetadataLimits::default() }; - let pipeline = PersistencePipeline::start(repository.clone(), 1, limits); + let pipeline = PersistencePipeline::start(repository.clone(), 1, limits, disk_guard()); let mut invalid = torrent(); invalid.files.push(FileInfo { path: "second".into(), @@ -387,4 +414,22 @@ mod tests { assert_eq!(rejections.len(), 1); assert_eq!(rejections[0].reason, MetadataRejectionReason::TooManyFiles); } + + #[tokio::test] + async fn read_only_guard_rejects_new_metadata_without_blocking_shutdown() { + let repository = Arc::new(MemoryRepository::default()); + let guard = DiskGuard::new(&DiskGuardConfig { + enabled: true, + check_interval_secs: 1, + minimum_free_bytes: 100, + resume_free_bytes: 200, + }); + guard.observe_available(99, 0); + let pipeline = + PersistencePipeline::start(repository.clone(), 1, MetadataLimits::default(), guard); + + assert!(!pipeline.ingress.try_enqueue(torrent())); + pipeline.close_and_join().await.unwrap(); + assert!(repository.records.lock().unwrap().is_empty()); + } } diff --git a/dht-search/src/disk_guard.rs b/dht-search/src/disk_guard.rs new file mode 100644 index 0000000..966452f --- /dev/null +++ b/dht-search/src/disk_guard.rs @@ -0,0 +1,346 @@ +// 负责监测数据盘剩余空间并协调所有后台写入进入可恢复的只读状态 + +use std::{ + io, + path::{Path, PathBuf}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::Duration, +}; + +use tokio_util::sync::CancellationToken; + +use crate::{config::DiskGuardConfig, crawler::pipeline::PersistenceIngress}; + +const UNKNOWN_AVAILABLE_BYTES: u64 = u64::MAX; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum DiskMode { + Normal, + Draining, + ReadOnly, +} + +impl DiskMode { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Normal => "normal", + Self::Draining => "draining", + Self::ReadOnly => "read_only", + } + } +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct DiskGuardSnapshot { + pub(crate) mode: DiskMode, + pub(crate) available_bytes: Option, + pub(crate) minimum_free_bytes: u64, + pub(crate) resume_free_bytes: u64, + pub(crate) active_writes: usize, + pub(crate) probe_failed: bool, + pub(crate) probe_failures: u64, + pub(crate) transitions: u64, + pub(crate) rejected_new_work: u64, +} + +#[derive(Clone)] +pub(crate) struct DiskGuard { + inner: Arc, +} + +struct DiskGuardInner { + enabled: bool, + minimum_free_bytes: u64, + resume_free_bytes: u64, + state: Mutex, + available_bytes: AtomicU64, + probe_failed: AtomicBool, + probe_failures: AtomicU64, + transitions: AtomicU64, + rejected_new_work: AtomicU64, +} + +struct GateState { + mode: DiskMode, + active_writes: usize, +} + +pub(crate) struct DiskWritePermit { + inner: Arc, +} + +pub(crate) async fn run( + guard: DiskGuard, + data_dir: PathBuf, + config: DiskGuardConfig, + persistence: PersistenceIngress, + cancel: CancellationToken, +) { + let mut ticker = tokio::time::interval(Duration::from_secs(config.check_interval_secs)); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = cancel.cancelled() => break, + _ = ticker.tick() => { + guard.probe(&data_dir, persistence.snapshot().queue_depth); + } + } + } +} + +impl DiskGuard { + pub(crate) fn new(config: &DiskGuardConfig) -> Self { + Self { + inner: Arc::new(DiskGuardInner { + enabled: config.enabled, + minimum_free_bytes: config.minimum_free_bytes, + resume_free_bytes: config.resume_free_bytes, + state: Mutex::new(GateState { + mode: DiskMode::Normal, + active_writes: 0, + }), + available_bytes: AtomicU64::new(UNKNOWN_AVAILABLE_BYTES), + probe_failed: AtomicBool::new(false), + probe_failures: AtomicU64::new(0), + transitions: AtomicU64::new(0), + rejected_new_work: AtomicU64::new(0), + }), + } + } + + pub(crate) fn probe(&self, path: &Path, persistence_queue: usize) { + match fs2::available_space(path) { + Ok(available) => self.observe_available(available, persistence_queue), + Err(error) => self.observe_probe_error(&error, persistence_queue), + } + } + + pub(crate) fn begin_admission(&self) -> Option { + let permit = self.begin_write(false); + if permit.is_none() { + self.inner.rejected_new_work.fetch_add(1, Ordering::Relaxed); + } + permit + } + + pub(crate) fn begin_new_write(&self) -> Option { + self.begin_write(false) + } + + pub(crate) fn begin_drain_write(&self) -> Option { + self.begin_write(true) + } + + #[cfg(test)] + pub(crate) fn mode(&self) -> DiskMode { + self.inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .mode + } + + pub(crate) fn snapshot(&self) -> DiskGuardSnapshot { + let state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let available = self.inner.available_bytes.load(Ordering::Relaxed); + DiskGuardSnapshot { + mode: state.mode, + available_bytes: (available != UNKNOWN_AVAILABLE_BYTES).then_some(available), + minimum_free_bytes: self.inner.minimum_free_bytes, + resume_free_bytes: self.inner.resume_free_bytes, + active_writes: state.active_writes, + probe_failed: self.inner.probe_failed.load(Ordering::Relaxed), + probe_failures: self.inner.probe_failures.load(Ordering::Relaxed), + transitions: self.inner.transitions.load(Ordering::Relaxed), + rejected_new_work: self.inner.rejected_new_work.load(Ordering::Relaxed), + } + } + + fn begin_write(&self, allow_draining: bool) -> Option { + if !self.inner.enabled { + let mut state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.active_writes = state.active_writes.saturating_add(1); + return Some(DiskWritePermit { + inner: self.inner.clone(), + }); + } + let mut state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let allowed = + state.mode == DiskMode::Normal || (allow_draining && state.mode == DiskMode::Draining); + if !allowed { + return None; + } + state.active_writes = state.active_writes.saturating_add(1); + Some(DiskWritePermit { + inner: self.inner.clone(), + }) + } + + pub(crate) fn observe_available(&self, available: u64, persistence_queue: usize) { + self.inner + .available_bytes + .store(available, Ordering::Relaxed); + self.inner.probe_failed.store(false, Ordering::Relaxed); + if !self.inner.enabled { + return; + } + if available < self.inner.minimum_free_bytes { + self.transition_to(DiskMode::Draining, Some(available), None); + self.finish_draining(persistence_queue); + } else if available >= self.inner.resume_free_bytes { + self.transition_to(DiskMode::Normal, Some(available), None); + } else { + self.finish_draining(persistence_queue); + } + } + + fn observe_probe_error(&self, error: &io::Error, persistence_queue: usize) { + self.inner + .available_bytes + .store(UNKNOWN_AVAILABLE_BYTES, Ordering::Relaxed); + self.inner.probe_failed.store(true, Ordering::Relaxed); + self.inner.probe_failures.fetch_add(1, Ordering::Relaxed); + if !self.inner.enabled { + return; + } + self.transition_to(DiskMode::Draining, None, Some(error)); + self.finish_draining(persistence_queue); + } + + fn finish_draining(&self, persistence_queue: usize) { + if persistence_queue != 0 { + return; + } + let mut state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.mode == DiskMode::Draining && state.active_writes == 0 { + let old = state.mode; + state.mode = DiskMode::ReadOnly; + drop(state); + self.record_transition(old, DiskMode::ReadOnly, None, None); + } + } + + fn transition_to(&self, target: DiskMode, available: Option, error: Option<&io::Error>) { + let mut state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.mode == target + || (state.mode == DiskMode::ReadOnly && target == DiskMode::Draining) + { + return; + } + let old = state.mode; + state.mode = target; + drop(state); + self.record_transition(old, target, available, error); + } + + fn record_transition( + &self, + old: DiskMode, + new: DiskMode, + available: Option, + error: Option<&io::Error>, + ) { + self.inner.transitions.fetch_add(1, Ordering::Relaxed); + match new { + DiskMode::Normal => tracing::info!( + previous = old.as_str(), + available_bytes = available, + "磁盘空间恢复并重新接受写入" + ), + DiskMode::Draining => tracing::warn!( + previous = old.as_str(), + available_bytes = available, + error = error.map(ToString::to_string), + "磁盘空间不足并停止接收新任务" + ), + DiskMode::ReadOnly => { + tracing::warn!(previous = old.as_str(), "待写入任务已排空并进入只读保护") + } + } + } +} + +impl Drop for DiskWritePermit { + fn drop(&mut self) { + let mut state = self + .inner + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.active_writes = state.active_writes.saturating_sub(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn guard() -> DiskGuard { + DiskGuard::new(&DiskGuardConfig { + enabled: true, + check_interval_secs: 1, + minimum_free_bytes: 100, + resume_free_bytes: 200, + }) + } + + #[test] + fn low_space_drains_then_enters_read_only_and_recovers_with_hysteresis() { + let guard = guard(); + let permit = guard.begin_new_write().unwrap(); + guard.observe_available(99, 1); + assert_eq!(guard.mode(), DiskMode::Draining); + assert!(guard.begin_admission().is_none()); + guard.observe_available(150, 0); + assert_eq!(guard.mode(), DiskMode::Draining); + drop(permit); + guard.observe_available(150, 0); + assert_eq!(guard.mode(), DiskMode::ReadOnly); + assert!(guard.begin_new_write().is_none()); + guard.observe_available(199, 0); + assert_eq!(guard.mode(), DiskMode::ReadOnly); + guard.observe_available(200, 0); + assert_eq!(guard.mode(), DiskMode::Normal); + } + + #[test] + fn probe_failure_conservatively_enters_read_only() { + let guard = guard(); + guard.observe_probe_error(&io::Error::other("probe failed"), 0); + let snapshot = guard.snapshot(); + assert_eq!(snapshot.mode, DiskMode::ReadOnly); + assert!(snapshot.probe_failed); + assert_eq!(snapshot.probe_failures, 1); + } + + #[test] + fn draining_allows_only_existing_queue_writes() { + let guard = guard(); + guard.observe_available(1, 1); + assert!(guard.begin_new_write().is_none()); + assert!(guard.begin_drain_write().is_some()); + } +} diff --git a/dht-search/src/index_worker.rs b/dht-search/src/index_worker.rs index f711959..4fab9a8 100644 --- a/dht-search/src/index_worker.rs +++ b/dht-search/src/index_worker.rs @@ -2,25 +2,51 @@ use std::{sync::Arc, time::Duration}; -use dht_search::{search::SearchEngine, storage::RocksTorrentRepository}; +use dht_search::{ + search::SearchEngine, + storage::{RocksTorrentRepository, TorrentRepository}, +}; use tokio_util::sync::CancellationToken; +use crate::disk_guard::DiskGuard; + +pub(crate) struct IndexWorkerOptions { + pub(crate) batch_size: usize, + pub(crate) interval: Duration, + pub(crate) prepare_full_reindex: bool, +} + pub(crate) async fn run( repository: Arc, search: SearchEngine, - batch_size: usize, - interval: Duration, + options: IndexWorkerOptions, + disk_guard: DiskGuard, cancel: CancellationToken, fatal: tokio::sync::oneshot::Sender, ) -> Result<(), String> { - let mut ticker = tokio::time::interval(interval); + let mut prepare_full_reindex = options.prepare_full_reindex; + let mut ticker = tokio::time::interval(options.interval); ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut consecutive_retries = 0_u32; loop { tokio::select! { _ = cancel.cancelled() => break, _ = ticker.tick() => { - match index_one_batch(repository.clone(), search.clone(), batch_size).await { + let Some(_permit) = disk_guard.begin_new_write() else { + continue; + }; + if prepare_full_reindex { + let rebuild_repository = repository.clone(); + let records = tokio::task::spawn_blocking(move || { + rebuild_repository.prepare_full_reindex() + }) + .await + .map_err(|error| error.to_string())? + .map_err(|error| error.to_string())?; + tracing::info!(records, "检测到新搜索索引并准备全量重建"); + prepare_full_reindex = false; + } + match index_one_batch(repository.clone(), search.clone(), options.batch_size).await { Ok(count) => { consecutive_retries = 0; if count > 0 { @@ -45,14 +71,18 @@ pub(crate) async fn run( } } - drain_before_shutdown(repository, search, batch_size).await + drain_before_shutdown(repository, search, options.batch_size, disk_guard).await } async fn drain_before_shutdown( repository: Arc, search: SearchEngine, batch_size: usize, + disk_guard: DiskGuard, ) -> Result<(), String> { + let Some(_permit) = disk_guard.begin_new_write() else { + return Ok(()); + }; let mut retries = 0_u32; loop { match index_one_batch(repository.clone(), search.clone(), batch_size).await { diff --git a/dht-search/src/main.rs b/dht-search/src/main.rs index fd825d1..70327ae 100644 --- a/dht-search/src/main.rs +++ b/dht-search/src/main.rs @@ -6,6 +6,7 @@ mod api; mod app; mod config; mod crawler; +mod disk_guard; mod error; mod index_worker; mod monitor; diff --git a/dht-search/src/monitor.rs b/dht-search/src/monitor.rs index 30964a1..3e67cbd 100644 --- a/dht-search/src/monitor.rs +++ b/dht-search/src/monitor.rs @@ -6,11 +6,12 @@ use dht_crawler::DHTServer; use dht_search::domain::MetadataRejectionReason; use tokio_util::sync::CancellationToken; -use crate::crawler::pipeline::PersistenceIngress; +use crate::{crawler::pipeline::PersistenceIngress, disk_guard::DiskGuard}; pub(crate) async fn run( server: DHTServer, ingress: PersistenceIngress, + disk_guard: DiskGuard, interval_secs: u64, cancel: CancellationToken, ) { @@ -25,6 +26,7 @@ pub(crate) async fn run( let dht = server.runtime_stats().snapshot(); let observability = server.runtime_stats().observability_snapshot(); let storage = ingress.snapshot(); + let disk = disk_guard.snapshot(); let udp_tx_per_second = observability .udp_tx_packets .saturating_sub(previous_udp_tx) @@ -78,6 +80,11 @@ pub(crate) async fn run( persistence_invalid = storage.invalid, persistence_failed = storage.failed, persistence_queue = storage.queue_depth, + disk_state = disk.mode.as_str(), + disk_available_bytes = disk.available_bytes, + disk_active_writes = disk.active_writes, + disk_probe_failed = disk.probe_failed, + disk_rejected_new_work = disk.rejected_new_work, "运行状态" ) } diff --git a/dht-search/src/verification.rs b/dht-search/src/verification.rs index f6e8566..854f737 100644 --- a/dht-search/src/verification.rs +++ b/dht-search/src/verification.rs @@ -15,13 +15,19 @@ use dht_search::{ use tokio::task::{JoinHandle, JoinSet}; use tokio_util::sync::CancellationToken; -use crate::config::VerificationConfig; +#[cfg(test)] +use crate::config::DiskGuardConfig; +use crate::{ + config::VerificationConfig, + disk_guard::{DiskGuard, DiskWritePermit}, +}; #[derive(Clone)] pub(crate) struct VerificationIngress { repository: Arc, capacity: usize, stats: VerificationStats, + disk_guard: DiskGuard, } #[derive(Clone, Default)] @@ -59,6 +65,7 @@ pub(crate) fn start( repository: Arc, server: DHTServer, config: VerificationConfig, + disk_guard: DiskGuard, cancel: CancellationToken, ) -> (VerificationIngress, JoinHandle>) { let stats = VerificationStats::default(); @@ -66,9 +73,13 @@ pub(crate) fn start( repository.verification_queue_len().unwrap_or_default() as u64, Ordering::Relaxed, ); - let ingress = - VerificationIngress::new(repository.clone(), config.queue_capacity, stats.clone()); - let task = tokio::spawn(run(repository, server, config, stats, cancel)); + let ingress = VerificationIngress::new( + repository.clone(), + config.queue_capacity, + stats.clone(), + disk_guard.clone(), + ); + let task = tokio::spawn(run(repository, server, config, stats, disk_guard, cancel)); (ingress, task) } @@ -77,26 +88,40 @@ impl VerificationIngress { repository: Arc, capacity: usize, stats: VerificationStats, + disk_guard: DiskGuard, ) -> Self { Self { repository, capacity, stats, + disk_guard, } } #[cfg(test)] pub(crate) fn for_test(repository: Arc, capacity: usize) -> Self { - Self::new(repository, capacity, VerificationStats::default()) + Self::new( + repository, + capacity, + VerificationStats::default(), + DiskGuard::new(&DiskGuardConfig { + enabled: false, + ..DiskGuardConfig::default() + }), + ) } pub(crate) async fn enqueue(&self, hashes: Vec, priority: VerificationPriority) { if hashes.is_empty() { return; } + let Some(permit) = self.disk_guard.begin_admission() else { + return; + }; let repository = self.repository.clone(); let capacity = self.capacity; let result = tokio::task::spawn_blocking(move || { + let _permit = permit; let outcome = repository.enqueue_verification(&hashes, priority, unix_timestamp(), capacity)?; let queue_len = repository.verification_queue_len()?; @@ -153,6 +178,7 @@ async fn run( server: DHTServer, config: VerificationConfig, stats: VerificationStats, + disk_guard: DiskGuard, cancel: CancellationToken, ) -> Result<(), String> { let mut active = JoinSet::new(); @@ -168,14 +194,18 @@ async fn run( } _ = ticker.tick() => { while active.len() < config.max_active { + let Some(permit) = disk_guard.begin_new_write() else { break }; let claim_repository = repository.clone(); let lease_secs = config.lease_secs; let request = tokio::task::spawn_blocking(move || { - claim_repository.claim_verification(unix_timestamp(), lease_secs) + claim_repository + .claim_verification(unix_timestamp(), lease_secs) + .map(|request| (request, permit)) }) .await .map_err(|error| error.to_string())? .map_err(|error| error.to_string())?; + let (request, permit) = request; let Some(request) = request else { break }; stats.inner.started.fetch_add(1, Ordering::Relaxed); active.spawn(verify_one( @@ -184,6 +214,7 @@ async fn run( request.info_hash, config.max_peer_attempts, stats.clone(), + permit, )); } } @@ -201,6 +232,7 @@ async fn verify_one( info_hash: InfoHash, max_peer_attempts: usize, stats: VerificationStats, + _permit: DiskWritePermit, ) -> Result<(), String> { let peer_repository = repository.clone(); let stored_peers = tokio::task::spawn_blocking(move || peer_repository.get(info_hash)) diff --git a/web/src/App.vue b/web/src/App.vue index 4219cda..9760029 100644 --- a/web/src/App.vue +++ b/web/src/App.vue @@ -8,6 +8,7 @@ import { Button } from '@/components/ui/button' import { AppPagination } from '@/components/ui/pagination' import { AppSelect } from '@/components/ui/select' import { getStats, getTorrent, getVariants, search } from '@/lib/api' +import { formatBytes } from '@/lib/format' import type { SearchHit, SearchPage, SearchSort, ServiceStats, TorrentDetail, ContentVariants } from '@/types/api' const query = ref('') @@ -62,6 +63,11 @@ const sortOptions: ReadonlyArray<{ value: SearchSort; label: string }> = [ { value: 'size_desc', label: '大小降序' }, { value: 'size_asc', label: '大小升序' }, ] +function diskStateLabel(state: ServiceStats['disk_state']): string { + if (state === 'normal') return '正常' + if (state === 'draining') return '正在排空' + return '只读保护' +} function updateBrowserUrl() { const params = new URLSearchParams() if (submittedQuery.value) params.set('q', submittedQuery.value) @@ -212,10 +218,10 @@ onBeforeUnmount(() => {

DHT Search

- +

服务运行状态

-
DHT 节点{{ stats.nodes.toLocaleString() }}
已索引内容{{ stats.indexed_documents.toLocaleString() }}
获取成功{{ stats.metadata_ok.toLocaleString() }}
下载中{{ stats.metadata_in_flight }}
新收录{{ stats.persistence_inserted.toLocaleString() }}
已过滤{{ stats.metadata_filtered.toLocaleString() }}
验证成功{{ stats.verification_succeeded.toLocaleString() }}
+
磁盘状态{{ diskStateLabel(stats.disk_state) }}
磁盘剩余{{ stats.disk_available_bytes === null ? '未知' : formatBytes(stats.disk_available_bytes) }}
DHT 节点{{ stats.nodes.toLocaleString() }}
已索引内容{{ stats.indexed_documents.toLocaleString() }}
获取成功{{ stats.metadata_ok.toLocaleString() }}
下载中{{ stats.metadata_in_flight }}
新收录{{ stats.persistence_inserted.toLocaleString() }}
已过滤{{ stats.metadata_filtered.toLocaleString() }}
验证成功{{ stats.verification_succeeded.toLocaleString() }}

无法获取服务状态

diff --git a/web/src/types/api.ts b/web/src/types/api.ts index 4f617d4..20d5105 100644 --- a/web/src/types/api.ts +++ b/web/src/types/api.ts @@ -95,6 +95,15 @@ export interface ContentVariants { } export interface ServiceStats { + disk_state: 'normal' | 'draining' | 'read_only' + disk_available_bytes: number | null + disk_minimum_free_bytes: number + disk_resume_free_bytes: number + disk_active_writes: number + disk_probe_failed: boolean + disk_probe_failures: number + disk_transitions: number + disk_rejected_new_work: number nodes: number metadata_ok: number metadata_failed: number