use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Instant, UNIX_EPOCH};
use anyhow::{Context, Result};
use tantivy::{Index, IndexWriter, doc};
use walkdir::WalkDir;
use crate::cursor::{UsnCursor, VolumeKey};
use crate::extract::extract_text;
use crate::meta::{IndexMeta, SCHEMA_VERSION, save_meta};
use crate::ocr;
use crate::ocr_queue::OcrQueue;
use crate::volume::{self, RootCapability};
use crate::{Fields, build_schema, register_tokenizers};
pub struct IndexStats {
pub indexed: usize,
pub skipped: usize,
pub seconds: f64,
}
pub(crate) const PROGRESS_INTERVAL: usize = 50;
#[derive(Debug, Clone)]
pub struct IndexProgress {
pub processed: usize,
pub path: PathBuf,
}
const SKIP_DIRS: &[&str] = &["node_modules", "target", ".git", ".venv", "__pycache__"];
const REMOVE_DIR_ALL_RETRIES: u32 = 4;
const REMOVE_DIR_ALL_RETRY_BASE_DELAY: std::time::Duration = std::time::Duration::from_millis(50);
pub fn remove_dir_all_retrying(path: &Path) -> std::io::Result<()> {
let mut delay = REMOVE_DIR_ALL_RETRY_BASE_DELAY;
for attempt in 1..=REMOVE_DIR_ALL_RETRIES {
match std::fs::remove_dir_all(path) {
Ok(()) => return Ok(()),
Err(err) if err.kind() == std::io::ErrorKind::PermissionDenied => {
eprintln!(
"删除目录 {} 遇到句柄未释放(第 {attempt}/{REMOVE_DIR_ALL_RETRIES} 次重试前等待 {delay:?}): {err}",
path.display()
);
std::thread::sleep(delay);
delay *= 2;
}
Err(err) => return Err(err),
}
}
std::fs::remove_dir_all(path)
}
pub(crate) fn walk_index_files(root: &Path) -> impl Iterator<Item = PathBuf> {
WalkDir::new(root)
.into_iter()
.filter_entry(|e| {
if e.depth() == 0 {
return true;
}
let name = e.file_name().to_string_lossy();
!(e.file_type().is_dir()
&& (SKIP_DIRS.contains(&name.as_ref()) || name.starts_with('.')))
})
.filter_map(|e| e.ok())
.filter(|e| e.file_type().is_file())
.map(|e| e.path().to_path_buf())
}
fn collect_index_files(target_dir: &Path) -> (Vec<PathBuf>, HashMap<VolumeKey, UsnCursor>) {
if let RootCapability::Fast { volume } = volume::probe_root_capability(target_dir)
&& let Some(result) = collect_via_mft(&volume, target_dir)
{
return result;
}
(walk_index_files(target_dir).collect(), HashMap::new())
}
#[cfg(windows)]
fn collect_via_mft(
volume: &VolumeKey,
target_dir: &Path,
) -> Option<(Vec<PathBuf>, HashMap<VolumeKey, UsnCursor>)> {
let cursor = match crate::usn::snapshot_journal(volume) {
Ok(cursor) => cursor,
Err(err) => {
eprintln!("拍 USN 游标快照失败({volume}),MFT 快速路径整体回退: {err}");
return None;
}
};
let (_table, files, stats) =
match crate::mft::enumerate(volume, std::slice::from_ref(&target_dir.to_path_buf())) {
Ok(result) => result,
Err(err) => {
eprintln!("MFT 快速枚举失败({volume}),回退到目录遍历: {err}");
return None;
}
};
eprintln!(
"MFT 快速枚举({volume}):整卷扫到 {} 条记录,落在监听根内 {} 条",
stats.scanned, stats.matched
);
let mut cursors = HashMap::new();
cursors.insert(volume.clone(), cursor);
Some((files, cursors))
}
#[cfg(not(windows))]
fn collect_via_mft(
_volume: &VolumeKey,
_target_dir: &Path,
) -> Option<(Vec<PathBuf>, HashMap<VolumeKey, UsnCursor>)> {
None
}
pub(crate) fn file_stat(path: &Path) -> Option<(i64, u64)> {
let meta = std::fs::metadata(path).ok()?;
let mtime = meta
.modified()
.ok()
.and_then(|t| t.duration_since(UNIX_EPOCH).ok())
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
Some((mtime, meta.len()))
}
pub(crate) fn add_file_document(
writer: &IndexWriter,
fields: &Fields,
path: &Path,
index_dir: &Path,
) -> Result<bool> {
if ocr::is_image(path) {
return add_image_document(writer, fields, path, index_dir);
}
let Some(content) = extract_text(path) else {
return Ok(false);
};
let (mtime, size) = file_stat(path).unwrap_or((0, 0));
let name = path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_default();
let ext = path
.extension()
.map(|e| e.to_string_lossy().to_ascii_lowercase())
.unwrap_or_default();
writer.add_document(doc!(
fields.path => path.to_string_lossy().into_owned(),
fields.name => name,
fields.ext => ext,
fields.content => content,
fields.mtime => mtime,
fields.size => size,
fields.kind => "text",
))?;
Ok(true)
}
fn add_image_document(
writer: &IndexWriter,
fields: &Fields,
path: &Path,
index_dir: &Path,
) -> Result<bool> {
let Some((mtime, size)) = file_stat(path) else {
return Ok(false);
};
if size > ocr::MAX_IMAGE_BYTES {
return Ok(false);
}
let queue = OcrQueue::for_index_dir(index_dir);
let content = match queue.cached_content(path, mtime, size) {
Some(cached) => cached,
None => {
queue.enqueue(path.to_path_buf(), mtime, size);
String::new()
}
};
add_image_document_with_content(writer, fields, path, mtime, size, &content)?;
Ok(true)
}
pub(crate) fn add_image_document_with_content(
writer: &IndexWriter,
fields: &Fields,
path: &Path,
mtime: i64,
size: u64,
content: &str,
) -> Result<()> {
let name = path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_default();
let ext = path
.extension()
.map(|e| e.to_string_lossy().to_ascii_lowercase())
.unwrap_or_default();
writer.add_document(doc!(
fields.path => path.to_string_lossy().into_owned(),
fields.name => name,
fields.ext => ext,
fields.kind => "image",
fields.content => content.to_string(),
fields.mtime => mtime,
fields.size => size,
))?;
Ok(())
}
pub(crate) fn commit_index_tail(
save_ocr_queue: impl FnOnce() -> Result<()>,
commit_writer: impl FnOnce() -> Result<()>,
) -> Result<()> {
save_ocr_queue()?;
commit_writer()
}
pub fn rebuild_index(index_dir: &Path, target_dir: &Path) -> Result<IndexStats> {
rebuild_index_with_progress(index_dir, target_dir, |_| {})
}
const REBUILD_RETRIES: u32 = 10;
const REBUILD_RETRY_BASE_DELAY: std::time::Duration = std::time::Duration::from_millis(300);
const REBUILD_RETRY_MAX_DELAY: std::time::Duration = std::time::Duration::from_secs(3);
fn writer_threads_for_attempt(attempt: u32) -> usize {
match attempt {
1..=4 => 4,
5..=7 => 2,
_ => 1,
}
}
pub(crate) fn is_transient_writer_killed(err: &anyhow::Error) -> bool {
err.chain().any(
|cause| match cause.downcast_ref::<tantivy::TantivyError>() {
Some(tantivy::TantivyError::ErrorInThread(_)) => true,
Some(tantivy::TantivyError::OpenWriteError(
tantivy::directory::error::OpenWriteError::IoError { io_error, .. },
)) => io_error.kind() == std::io::ErrorKind::PermissionDenied,
Some(tantivy::TantivyError::IoError(io_error)) => {
io_error.kind() == std::io::ErrorKind::PermissionDenied
}
_ => false,
},
)
}
pub fn rebuild_index_with_progress(
index_dir: &Path,
target_dir: &Path,
mut on_progress: impl FnMut(IndexProgress),
) -> Result<IndexStats> {
let start = Instant::now();
let mut delay = REBUILD_RETRY_BASE_DELAY;
for attempt in 1..=REBUILD_RETRIES {
let writer_threads = writer_threads_for_attempt(attempt);
match rebuild_index_attempt(index_dir, target_dir, writer_threads, &mut on_progress) {
Ok((indexed, skipped, usn_cursors)) => {
return finish_rebuild(index_dir, target_dir, indexed, skipped, usn_cursors, start);
}
Err(err) if attempt < REBUILD_RETRIES && is_transient_writer_killed(&err) => {
eprintln!(
"全量重建撞上瞬时的索引写入中断(第 {attempt}/{REBUILD_RETRIES} 次,常见于杀毒/EDR \
软件实时扫描新建文件),等待 {delay:?} 后整次重试: {err}"
);
std::thread::sleep(delay);
delay = (delay * 2).min(REBUILD_RETRY_MAX_DELAY);
continue;
}
Err(err) => return Err(err),
}
}
unreachable!("循环要么在 Ok 分支返回,要么在最后一次尝试的 Err 分支返回")
}
fn rebuild_index_attempt(
index_dir: &Path,
target_dir: &Path,
writer_threads: usize,
on_progress: &mut impl FnMut(IndexProgress),
) -> Result<(usize, usize, HashMap<VolumeKey, UsnCursor>)> {
if index_dir.exists() {
remove_dir_all_retrying(index_dir).context("清理旧索引目录失败")?;
}
std::fs::create_dir_all(index_dir)?;
let (schema, fields) = build_schema();
let index = Index::create_in_dir(index_dir, schema)?;
register_tokenizers(&index);
let mut writer: IndexWriter =
index.writer_with_num_threads(writer_threads, 200 * 1024 * 1024)?;
let mut indexed = 0usize;
let mut skipped = 0usize;
let (files, usn_cursors) = collect_index_files(target_dir);
for path in files {
if add_file_document(&writer, &fields, &path, index_dir)? {
indexed += 1;
} else {
skipped += 1;
}
let processed = indexed + skipped;
if processed % PROGRESS_INTERVAL == 0 {
on_progress(IndexProgress {
processed,
path: path.clone(),
});
}
}
let ocr_queue = OcrQueue::for_index_dir(index_dir);
ocr_queue.compact(&[target_dir.to_path_buf()]);
commit_index_tail(
|| ocr_queue.save().context("保存 OCR 队列状态失败"),
|| writer.commit().map(|_| ()).context("索引提交失败"),
)?;
writer
.wait_merging_threads()
.context("等待索引合并线程退出失败")?;
Ok((indexed, skipped, usn_cursors))
}
fn finish_rebuild(
index_dir: &Path,
target_dir: &Path,
indexed: usize,
skipped: usize,
usn_cursors: HashMap<VolumeKey, UsnCursor>,
start: Instant,
) -> Result<IndexStats> {
save_meta(
index_dir,
&IndexMeta {
schema_version: SCHEMA_VERSION,
roots: vec![target_dir.to_path_buf()],
usn_cursors,
},
)?;
Ok(IndexStats {
indexed,
skipped,
seconds: start.elapsed().as_secs_f64(),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rebuild_index_root_dot_prefixed_dir_is_not_skipped() -> Result<()> {
let index_dir = tempfile::tempdir()?;
let target_dir = tempfile::Builder::new().prefix(".dowse-test-").tempdir()?;
std::fs::write(target_dir.path().join("note.txt"), "hello dowse")?;
let stats = rebuild_index(index_dir.path(), target_dir.path())?;
assert_eq!(stats.indexed, 1);
Ok(())
}
#[test]
fn rebuild_index_with_progress_reports_every_interval() -> Result<()> {
let index_dir = tempfile::tempdir()?;
let target_dir = tempfile::tempdir()?;
let total = PROGRESS_INTERVAL * 2 + 3;
for i in 0..total {
std::fs::write(target_dir.path().join(format!("f{i}.txt")), "内容")?;
}
let mut reports = Vec::new();
let stats = rebuild_index_with_progress(index_dir.path(), target_dir.path(), |p| {
reports.push(p.processed);
})?;
assert_eq!(stats.indexed, total);
assert_eq!(
reports,
vec![PROGRESS_INTERVAL, PROGRESS_INTERVAL * 2],
"总数不是间隔整数倍时,最后一段不足一个间隔的尾巴不应触发额外回调"
);
Ok(())
}
#[test]
fn rebuild_index_with_progress_below_interval_reports_nothing() -> Result<()> {
let index_dir = tempfile::tempdir()?;
let target_dir = tempfile::tempdir()?;
std::fs::write(target_dir.path().join("only.txt"), "内容")?;
let mut calls = 0usize;
rebuild_index_with_progress(index_dir.path(), target_dir.path(), |_| calls += 1)?;
assert_eq!(calls, 0);
Ok(())
}
#[test]
fn commit_index_tail_saves_ocr_queue_before_committing_writer() {
let order = std::cell::RefCell::new(Vec::new());
let result = commit_index_tail(
|| {
order.borrow_mut().push("save_queue");
Ok(())
},
|| {
order.borrow_mut().push("commit_writer");
Ok(())
},
);
assert!(result.is_ok());
assert_eq!(order.into_inner(), vec!["save_queue", "commit_writer"]);
}
#[test]
fn commit_index_tail_skips_commit_when_queue_save_fails() {
let order = std::cell::RefCell::new(Vec::new());
let result = commit_index_tail(
|| {
order.borrow_mut().push("save_queue");
anyhow::bail!("模拟保存失败")
},
|| {
order.borrow_mut().push("commit_writer");
Ok(())
},
);
assert!(result.is_err());
assert_eq!(
order.into_inner(),
vec!["save_queue"],
"队列保存失败时不该继续提交索引写入"
);
}
#[test]
fn remove_dir_all_retrying_removes_existing_directory() -> Result<()> {
let dir = tempfile::tempdir()?;
let target = dir.path().join("victim");
std::fs::create_dir_all(&target)?;
std::fs::write(target.join("f.txt"), "内容")?;
remove_dir_all_retrying(&target)?;
assert!(!target.exists());
Ok(())
}
#[test]
fn remove_dir_all_retrying_propagates_non_permission_errors() {
let dir = tempfile::tempdir().unwrap();
let missing = dir.path().join("does-not-exist");
let err =
remove_dir_all_retrying(&missing).expect_err("目录不存在应该报错,不应该重试掩盖");
assert_eq!(err.kind(), std::io::ErrorKind::NotFound);
}
}