From 4c16c68921990aff3f91c1a11d467e9b7a3a23ce Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 12 Oct 2025 17:15:34 +0800 Subject: [PATCH] Add multiple jobs support for import task --- Cargo.toml | 11 +++--- src/args.rs | 9 +++-- src/main.rs | 90 ++++++++++++++++++++++++++++++++++++++---------- src/types.rs | 1 - src/utils/mod.rs | 1 - 5 files changed, 82 insertions(+), 30 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 45b1aa6..f1b765b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,7 +33,7 @@ markup5ever_rcdom = { version = "0.35", optional = true } memchr = { version = "2.7", optional = true } mozjpeg = { version = "0.10", optional = true } msg_tool_macro = { version = "0.2.9" } -num_cpus = { version = "1.17", optional = true } +num_cpus = "1.17" overf = "0.1" parse-size = { version = "1.1", optional = true } pelite = { version = "0.10", optional = true } @@ -67,7 +67,7 @@ artemis-arc = ["artemis", "msg_tool_macro/artemis-arc", "sha1"] bgi = ["fancy-regex"] bgi-arc = ["bgi", "rand", "utils-bit-stream"] bgi-audio = ["bgi"] -bgi-img = ["bgi", "image", "rand", "utils-threadpool", "utils-bit-stream"] +bgi-img = ["bgi", "image", "rand", "utils-bit-stream"] cat-system = ["fancy-regex", "flate2", "int-enum"] cat-system-arc = ["cat-system", "pelite", "utils-blowfish", "utils-crc32"] cat-system-img = ["cat-system", "flate2", "image", "mozjpeg", "utils-bit-stream"] @@ -86,7 +86,7 @@ hexen-haus = ["memchr", "utils-str"] hexen-haus-arc = ["hexen-haus"] hexen-haus-img = ["hexen-haus", "image"] kirikiri = ["emote-psb", "fancy-regex", "flate2", "json", "lz4", "utils-escape"] -kirikiri-arc = ["kirikiri", "adler", "fastcdc", "flate2", "parse-size", "sha2", "xp3", "zstd", "utils-threadpool"] +kirikiri-arc = ["kirikiri", "adler", "fastcdc", "flate2", "parse-size", "sha2", "xp3", "zstd"] kirikiri-img = ["kirikiri", "image", "libtlg-rs"] silky = [] softpal = ["int-enum"] @@ -97,9 +97,9 @@ will-plus-img = ["will-plus", "image"] yaneurao = [] yaneurao-itufuru = ["yaneurao"] # basic feature -image = ["png", "utils-threadpool"] +image = ["png"] image-jpg = ["mozjpeg"] -image-jxl = ["image", "jpegxl-sys", "utils-threadpool"] +image-jxl = ["image", "jpegxl-sys"] image-webp = ["webp"] lossless-audio = ["utils-pcm"] audio-flac = ["libflac-sys", "utils-pcm"] @@ -112,7 +112,6 @@ utils-crc32 = [] utils-escape = ["fancy-regex"] utils-pcm = [] utils-str = [] -utils-threadpool = ["num_cpus"] [target.'cfg(windows)'.dependencies] windows-sys = { version = "0.61", features = ["Win32_Globalization", "Win32_System_Diagnostics_Debug"] } diff --git a/src/args.rs b/src/args.rs index c1acf7a..be7810f 100644 --- a/src/args.rs +++ b/src/args.rs @@ -72,7 +72,7 @@ fn parse_jxl_distance(s: &str) -> Result { } /// Tools for export and import scripts -#[derive(Parser, Debug)] +#[derive(Parser, Debug, Clone)] #[clap( group = ArgGroup::new("encodingg").multiple(false), group = ArgGroup::new("output_encodingg").multiple(false), @@ -530,7 +530,7 @@ pub struct Arg { pub command: Command, } -#[derive(Parser, Debug)] +#[derive(Parser, Debug, Clone)] #[clap(group = ArgGroup::new("patched_encodingg").multiple(false), group = ArgGroup::new("patched_archive_encodingg").multiple(false))] pub struct ImportArgs { /// Input script file or directory @@ -593,9 +593,12 @@ pub struct ImportArgs { pub replacement_json: Option, #[arg(long, action = ArgAction::SetTrue)] pub warn_when_output_file_not_found: bool, + #[arg(short = 'j', long, default_value_t = 1)] + /// Workers count for import scripts in parallel. + pub jobs: usize, } -#[derive(Subcommand, Debug)] +#[derive(Subcommand, Debug, Clone)] /// Commands pub enum Command { /// Extract from script diff --git a/src/main.rs b/src/main.rs index 4f94735..30979e6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2603,6 +2603,7 @@ fn main() { std::process::exit(1); }); let arg = args::parse_args(); + let argn = std::sync::Arc::new(arg.clone()); if arg.backtrace { unsafe { std::env::set_var("RUST_LIB_BACKTRACE", "1") }; } @@ -2870,7 +2871,7 @@ fn main() { } None => None, }; - let repl = match &args.replacement_json { + let repl = std::sync::Arc::new(match &args.replacement_json { Some(replacement_json) => { let b = utils::files::read_file(replacement_json).unwrap(); let s = String::from_utf8(b).unwrap(); @@ -2878,7 +2879,7 @@ fn main() { Some(table) } None => None, - }; + }); let (scripts, is_dir) = utils::files::collect_files(&args.input, arg.recursive, false).unwrap(); if is_dir { @@ -2897,25 +2898,76 @@ fn main() { } else { None }; + let workers = if args.jobs > 1 { + Some( + utils::threadpool::ThreadPool::<()>::new( + args.jobs, + Some("import-worker-"), + true, + ) + .unwrap(), + ) + } else { + None + }; for script in scripts.iter() { - let re = import_script( - &script, - &arg, - cfg.clone(), - args, - root_dir, - name_csv.as_ref(), - repl.as_ref(), - ); - match re { - Ok(s) => { - COUNTER.inc(s); - } - Err(e) => { + if let Some(workers) = workers.as_ref() { + let arg = argn.clone(); + let cfg = cfg.clone(); + let script = script.clone(); + let name_csv = name_csv.as_ref().map(|s| s.clone()); + let repl = repl.clone(); + let root_dir = root_dir.map(|s| s.to_path_buf()); + let args = args.clone(); + if let Err(e) = workers.execute( + move |_| { + let re = import_script( + &script, + &arg, + cfg, + &args, + root_dir.as_ref().map(|s| s.as_path()), + name_csv.as_ref(), + (*repl).as_ref(), + ); + match re { + Ok(s) => { + COUNTER.inc(s); + } + Err(e) => { + COUNTER.inc_error(); + eprintln!("Error exporting {}: {}", script, e); + if arg.backtrace { + eprintln!("Backtrace: {}", e.backtrace()); + } + } + } + }, + true, + ) { COUNTER.inc_error(); - eprintln!("Error exporting {}: {}", script, e); - if arg.backtrace { - eprintln!("Backtrace: {}", e.backtrace()); + eprintln!("Error executing import worker: {}", e); + } + } else { + let re = import_script( + &script, + &arg, + cfg.clone(), + args, + root_dir, + name_csv.as_ref(), + (*repl).as_ref(), + ); + match re { + Ok(s) => { + COUNTER.inc(s); + } + Err(e) => { + COUNTER.inc_error(); + eprintln!("Error exporting {}: {}", script, e); + if arg.backtrace { + eprintln!("Backtrace: {}", e.backtrace()); + } } } } diff --git a/src/types.rs b/src/types.rs index 005d836..e12caaa 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1063,7 +1063,6 @@ impl AsRef for LosslessAudioFormat { } } -#[cfg(feature = "utils-threadpool")] #[allow(unused)] pub(crate) fn get_default_threads() -> usize { num_cpus::get().max(2) / 2 diff --git a/src/utils/mod.rs b/src/utils/mod.rs index deeebf1..0381bfc 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -28,7 +28,6 @@ pub mod pcm; #[cfg(feature = "utils-str")] pub mod str; pub mod struct_pack; -#[cfg(feature = "utils-threadpool")] pub mod threadpool; #[cfg(windows)]