use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use humansize::{format_size, BINARY};
use kdam::{tqdm, Bar, BarExt};
use log::{debug, info, warn};
use tokio::sync::mpsc::Sender;
use tokio::sync::{Mutex, Semaphore};
use tokio::task::JoinSet;
use crate::hashing::{hash_file, hash_file_chunk, hash_mmap_file};
use crate::traversal::traverse_paths;
use crate::types::{DupeParams, DupeSet, Hash, PathCheckMode, PathGroup, SizeMap};
#[derive(Clone)]
pub struct DupeFinder {
pub min_file_size: u64,
pub read_semaphore: Arc<Semaphore>,
pub dupe_files: Arc<AtomicU64>,
pub dupe_sizes: Arc<AtomicU64>,
pub read_concurrency: usize,
pub disable_mmap: bool,
}
pub struct DupeProgress {
pbar_count: Bar,
pbar_size: Bar,
}
impl DupeProgress {
pub fn new(count: usize, total_size: u64, desc: String) -> Self {
let pbar_count = tqdm!(
total = count,
desc = desc.clone(),
unit = " file",
position = 0
);
let pbar_size = tqdm!(
total = total_size as usize,
desc = desc,
position = 1,
unit = "B",
unit_scale = true
);
Self {
pbar_count,
pbar_size,
}
}
pub fn update(&mut self, size: usize) {
let _ = self.pbar_count.update(1);
let _ = self.pbar_size.update(size);
}
}
impl DupeFinder {
pub fn new(params: DupeParams) -> Self {
Self {
min_file_size: params.min_file_size,
read_semaphore: Arc::new(Semaphore::new(params.read_concurrency)),
read_concurrency: params.read_concurrency,
dupe_files: Arc::default(),
dupe_sizes: Arc::default(),
disable_mmap: params.disable_mmap,
}
}
async fn process_wave(
&self,
tasks: Vec<(u64, usize, PathGroup)>, mode_fn: impl Fn(u64) -> PathCheckMode,
desc: &str,
) -> HashMap<(u64, usize, Hash), Vec<PathGroup>> {
let count = tasks.len();
let total_size = tasks.iter().map(|(s, _, _)| *s).sum();
let pbar = Arc::new(Mutex::new(DupeProgress::new(
count,
total_size,
desc.to_string(),
)));
let mut results: HashMap<(u64, usize, Hash), Vec<PathGroup>> = HashMap::new();
let mut set = JoinSet::new();
for (fsize, group_id, group) in tasks {
let permit = self.read_semaphore.clone().acquire_owned().await.unwrap();
let df = self.clone();
let mode = mode_fn(fsize);
let path = group.paths[0].clone();
set.spawn(async move {
let _permit = permit;
let hash = match mode {
PathCheckMode::Start => hash_file_chunk(&path, true, fsize, df.min_file_size).await,
PathCheckMode::End => hash_file_chunk(&path, false, fsize, df.min_file_size).await,
PathCheckMode::Full => {
tokio::task::spawn_blocking(move || {
if df.disable_mmap {
hash_file(&path)
} else {
hash_mmap_file(&path)
}
})
.await
.unwrap()
}
};
(fsize, group_id, group, hash)
});
}
while let Some(res) = set.join_next().await {
let (fsize, group_id, group, hash) = res.unwrap();
pbar.lock().await.update(fsize as usize);
match hash {
Ok(hash) => {
results.entry((fsize, group_id, hash)).or_default().push(group);
}
Err(e) => warn!("Failed to hash {}: {}", group.paths[0], e),
}
}
drop(pbar);
eprintln!();
eprintln!();
let dupes = results.values().filter(|v| v.len() > 1).count();
info!(
"{}: Processed {} files, found {} candidate groups",
desc, count, dupes
);
results
}
pub async fn traverse_paths(&self, paths: Vec<String>) -> Result<SizeMap, std::io::Error> {
traverse_paths(paths, self.min_file_size).await
}
pub fn find_duplicates(&self, size_map: SizeMap) -> tokio::sync::mpsc::Receiver<DupeSet> {
let (tx, rx) = tokio::sync::mpsc::channel(32);
let df = self.clone();
tokio::spawn(async move {
let _ = df.check_hashes_and_content(size_map, tx).await;
});
rx
}
pub async fn check_hashes_and_content(
&self,
size_map: SizeMap,
dupes_tx: Sender<DupeSet>,
) -> std::io::Result<()> {
let mut all_tasks: Vec<(u64, usize, PathGroup)> = Vec::new();
for (fsize, groups) in size_map {
for group in groups {
all_tasks.push((fsize, 0, group));
}
}
all_tasks.sort_by_key(|(_, _, group)| group.inode);
let mut wave1_results = self
.process_wave(
all_tasks,
|fsize| {
if fsize <= crate::hashing::EDGE_SIZE as u64 {
PathCheckMode::Full
} else {
PathCheckMode::Start
}
},
"Wave 1/3 (Initial)",
)
.await;
let mut final_groups: Vec<(u64, Vec<PathGroup>)> = Vec::new();
let mut wave2_tasks: Vec<(u64, usize, PathGroup)> = Vec::new();
let mut next_group_id = 0;
for ((fsize, _, _hash), groups) in wave1_results.drain() {
if groups.len() > 1 {
if fsize <= crate::hashing::EDGE_SIZE as u64 {
final_groups.push((fsize, groups));
} else {
next_group_id += 1;
for g in groups {
wave2_tasks.push((fsize, next_group_id, g));
}
}
}
}
wave2_tasks.sort_by_key(|(_, _, group)| group.inode);
let mut wave2_results = self
.process_wave(
wave2_tasks,
|_| PathCheckMode::End,
"Wave 2/3 (End)",
)
.await;
let mut wave3_tasks: Vec<(u64, usize, PathGroup)> = Vec::new();
next_group_id = 0;
for ((fsize, _, _hash), groups) in wave2_results.drain() {
if groups.len() > 1 {
next_group_id += 1;
for g in groups {
wave3_tasks.push((fsize, next_group_id, g));
}
}
}
wave3_tasks.sort_by_key(|(_, _, group)| group.inode);
let mut wave3_results = self
.process_wave(
wave3_tasks,
|_| PathCheckMode::Full,
"Wave 3/3 (Full)",
)
.await;
for ((fsize, _, _hash), groups) in wave3_results.drain() {
if groups.len() > 1 {
final_groups.push((fsize, groups));
}
}
for (fsize, groups) in final_groups {
let num_groups = groups.len();
let all_paths: Vec<String> = groups.into_iter().flat_map(|g| g.paths).collect();
debug!("Dupes found: {:?}", all_paths);
self.dupe_files
.fetch_add((all_paths.len() - 1) as u64, Ordering::Relaxed);
self.dupe_sizes
.fetch_add((num_groups - 1) as u64 * fsize, Ordering::Relaxed);
let _ = dupes_tx
.send(DupeSet {
fsize,
paths: all_paths,
})
.await;
}
info!("Dupes found: {:?}", self.dupe_files.load(Ordering::Relaxed));
info!(
"Dupes total size: {} (redundant data)",
format_size(self.dupe_sizes.load(Ordering::Relaxed), BINARY)
);
Ok(())
}
}