From a0f157774473a4b29dde41973fdf10adcb59c7b7 Mon Sep 17 00:00:00 2001 From: unspeaker Date: Tue, 11 Mar 2025 15:24:15 +0200 Subject: [PATCH] multithreaded scan --- src/main.rs | 67 +++++++++++++++++++++++++++++++++++++++++++++++----- src/model.rs | 46 ++++++++++++------------------------ 2 files changed, 76 insertions(+), 37 deletions(-) diff --git a/src/main.rs b/src/main.rs index 6b95aac..e40ca87 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,8 @@ use std::sync::{Arc, RwLock}; use std::path::{Path, PathBuf}; use std::env::{current_dir, set_current_dir}; use std::fs::read; +use std::thread::{sleep, spawn, JoinHandle}; +use std::time::Duration; use tek_tui::*; use tek_tui::tek_output::*; @@ -12,7 +14,7 @@ use tek_tui::tek_input::*; use crate::crossterm::event::{Event, KeyEvent, KeyCode, KeyModifiers, KeyEventState, KeyEventKind}; use clap::{arg, command, value_parser}; -use walkdir::WalkDir; +use walkdir::{WalkDir, DirEntry}; use xxhash_rust::xxh3::xxh3_64; use file_type::FileType; @@ -29,9 +31,9 @@ fn cli () -> clap::Command { command!() .arg(arg!([path] "Path to root directory") .value_parser(value_parser!(PathBuf))) - //.arg(arg!(-j --threads "Number of indexing threads") - //.required(false) - //.value_parser(value_parser!(usize))) + .arg(arg!(-j --threads "Number of indexing threads") + .required(false) + .value_parser(value_parser!(usize))) } fn main () -> Usually<()> { @@ -41,7 +43,60 @@ fn main () -> Usually<()> { } else { current_dir()? }; + let threads = args.get_one::("threads").map(|x|*x).unwrap_or_default().max(1); set_current_dir(&path)?; - let state = Arc::new(RwLock::new(Taggart::new(&path)?)); - Tui::new()?.run(&state) + let results = collect(&path, threads)?; + if let Ok(results) = Arc::try_unwrap(results) { + let mut results = results.into_inner()?; + results.sort(); + let state = Arc::new(RwLock::new(Taggart::new(&path, results)?)); + return Tui::new()?.run(&state) + } else { + panic!("read did not finish") + } + Ok(()) +} + +fn collect (root: &impl AsRef, thread_count: usize) -> Usually>>> { + let results = Arc::new(RwLock::new(vec![])); + let entries = Arc::new(RwLock::new(WalkDir::new(&root).into_iter() + .filter_entry(|e|!e.file_name().to_str().map(|s|s.starts_with(".")).unwrap_or(false)))); + let mut threads: Vec> = vec![]; + for thread_id in (0..thread_count).map(|x|x+1) { + threads.push({ + let root = root.as_ref().to_path_buf().clone(); + let results = results.clone(); + let entries = entries.clone(); + spawn(move || loop { + let (path, depth) = { + let entry = entries.write().unwrap().next(); + if let Some(entry) = entry { + let entry = entry.expect("failed to walk entry"); + let path = entry.path().to_path_buf(); + let depth = entry.depth(); + (path, depth) + } else { + break + } + }; + if depth > 0 { + let short_path = path.strip_prefix(root.as_path()) + .expect("failed to strip prefix"); + println!("(thread {thread_id}) {}", short_path.display()); + if let Ok(Some(entry)) = Entry::new(&root, &path, depth) { + results.write().unwrap().push(entry); + } + } + }) + }); + } + let timer = Duration::from_millis(100); + spawn(move || loop { + if threads.iter().all(|x|x.is_finished()) { + break + } else { + sleep(timer) + } + }).join(); + Ok(results) } diff --git a/src/model.rs b/src/model.rs index 6b84b56..336a609 100644 --- a/src/model.rs +++ b/src/model.rs @@ -57,34 +57,18 @@ pub enum EntryInfo { } impl Taggart { - pub fn new (root: &impl AsRef) -> Usually { + pub fn new (root: &impl AsRef, paths: Vec) -> Usually { Ok(Self { _root: root.as_ref().into(), - paths: Self::collect(root)?, cursor: 0, offset: 0, column: 0, size: Measure::new(), editing: None, columns: Columns::default(), + paths, }) } - pub fn collect (root: &impl AsRef) -> Usually> { - let mut paths = vec![]; - for entry in WalkDir::new(&root).into_iter() - .filter_entry(|e|!e.file_name().to_str().map(|s|s.starts_with(".")).unwrap_or(false)) - { - let entry = entry?; - if entry.depth() == 0 { - continue - } - if let Some(entry) = Entry::new(root, &entry)? { - paths.push(entry); - } - } - paths.sort(); - Ok(paths) - } pub fn edit_begin (&mut self) { let value = (self.columns.0[self.column].value)(&self.paths[self.cursor]); let value = format!("{}", value.unwrap_or_default()); @@ -110,20 +94,20 @@ impl Taggart { } impl Entry { - pub fn new (root: &impl AsRef, entry: &DirEntry) -> Perhaps { - println!("{}", entry.path().display()); - if entry.path().is_dir() { - Self::new_dir(root, entry) - } else if entry.path().is_file() { - Self::new_file(root, entry) + pub fn new (root: &impl AsRef, path: &impl AsRef, depth: usize) -> Perhaps { + let path = path.as_ref(); + if path.is_dir() { + Self::new_dir(root, &path, depth) + } else if path.is_file() { + Self::new_file(root, &path, depth) } else { Ok(None) } } - fn new_dir (root: &impl AsRef, entry: &DirEntry) -> Perhaps { + fn new_dir (root: &impl AsRef, path: &Path, depth: usize) -> Perhaps { Ok(Some(Self { - depth: entry.depth(), - path: entry.path().strip_prefix(root.as_ref())?.into(), + depth, + path: path.into(), info: EntryInfo::Directory { hash_file: None, catalog_file: None, @@ -132,11 +116,11 @@ impl Entry { }, })) } - fn new_file (root: &impl AsRef, entry: &DirEntry) -> Perhaps { + fn new_file (root: &impl AsRef, path: &Path, depth: usize) -> Perhaps { Ok(Some(Self { - depth: entry.depth(), - path: entry.path().strip_prefix(root.as_ref())?.into(), - info: EntryInfo::new(&read(entry.path())?)? + depth, + path: path.into(), + info: EntryInfo::new(&read(path)?)? })) } pub fn is_dir (&self) -> bool {