diff --git a/nod/Cargo.toml b/nod/Cargo.toml index a8e1a9a..f8f1d81 100644 --- a/nod/Cargo.toml +++ b/nod/Cargo.toml @@ -15,13 +15,14 @@ keywords.workspace = true categories = ["command-line-utilities", "parser-implementations"] [features] -default = ["compress-bzip2", "compress-lzma", "compress-zlib", "compress-zstd"] +default = ["compress-bzip2", "compress-lzma", "compress-zlib", "compress-zstd", "threading"] compress-bzip2 = ["bzip2"] compress-lzma = ["liblzma", "liblzma-sys"] compress-zlib = ["adler2", "miniz_oxide"] compress-zstd = ["zstd", "zstd-safe"] openssl = ["dep:openssl"] openssl-vendored = ["openssl", "openssl/vendored"] +threading = ["dep:crossbeam-channel", "dep:crossbeam-utils"] [dependencies] adler2 = { version = "2.0", optional = true } @@ -32,8 +33,8 @@ bytes = "1.10" bzip2 = { version = "0.6", features = ["static"], optional = true } cbc = "0.2.0-rc.1" crc32fast = "1.5" -crossbeam-channel = "0.5" -crossbeam-utils = "0.8" +crossbeam-channel = { version = "0.5", optional = true } +crossbeam-utils = { version = "0.8", optional = true } digest = { workspace = true } dyn-clone = "1.0" encoding_rs = "0.8" diff --git a/nod/src/disc/preloader.rs b/nod/src/disc/preloader.rs index 7520458..3230b06 100644 --- a/nod/src/disc/preloader.rs +++ b/nod/src/disc/preloader.rs @@ -1,20 +1,25 @@ +#[cfg(feature = "threading")] +use std::{collections::HashMap, thread::JoinHandle, time::Instant}; use std::{ - collections::HashMap, fmt::{Display, Formatter}, io, num::NonZeroUsize, sync::{Arc, Mutex}, - thread::JoinHandle, - time::{Duration, Instant}, + time::Duration, }; use bytes::{Bytes, BytesMut}; +#[cfg(feature = "threading")] use crossbeam_channel::{Receiver, Sender}; +#[cfg(feature = "threading")] use crossbeam_utils::sync::WaitGroup; use lru::LruCache; use polonius_the_crab::{polonius, polonius_return}; +#[cfg(feature = "threading")] use simple_moving_average::{SMA, SingleSumSMA}; -use tracing::{Level, debug, error, instrument, span}; +#[cfg(feature = "threading")] +use tracing::{Level, span}; +use tracing::{debug, error, instrument}; use zerocopy::FromZeros; use crate::{ @@ -59,6 +64,7 @@ pub struct SectorGroup { pub start_sector: u32, pub data: Bytes, pub sector_bitmap: u64, + #[allow(unused)] pub io_duration: Option, #[allow(unused)] // TODO WIA hash exceptions pub group_hashes: Option>, @@ -76,10 +82,15 @@ pub type SectorGroupResult = io::Result; #[allow(unused)] pub struct Preloader { + #[cfg(feature = "threading")] request_tx: Sender, + #[cfg(feature = "threading")] request_rx: Receiver, + #[cfg(feature = "threading")] stat_tx: Sender, + #[cfg(feature = "threading")] stat_rx: Receiver, + #[cfg(feature = "threading")] threads: Mutex, cache: Arc>, // Fallback single-threaded loader @@ -87,6 +98,7 @@ pub struct Preloader { } #[allow(unused)] +#[cfg(feature = "threading")] struct PreloaderThreads { join_handles: Vec>, last_adjust: Instant, @@ -96,6 +108,7 @@ struct PreloaderThreads { io_time_avg: SingleSumSMA, } +#[cfg(feature = "threading")] impl PreloaderThreads { fn new(join_handles: Vec>) -> Self { Self { @@ -153,6 +166,7 @@ impl PreloaderThreads { } struct PreloaderCache { + #[cfg(feature = "threading")] inflight: HashMap, lru_cache: LruCache, } @@ -160,6 +174,7 @@ struct PreloaderCache { impl Default for PreloaderCache { fn default() -> Self { Self { + #[cfg(feature = "threading")] inflight: Default::default(), lru_cache: LruCache::new(NonZeroUsize::new(64).unwrap()), } @@ -169,17 +184,21 @@ impl Default for PreloaderCache { impl PreloaderCache { fn push(&mut self, request: SectorGroupRequest, group: SectorGroup) { self.lru_cache.push(request, group); + #[cfg(feature = "threading")] self.inflight.remove(&request); } + #[cfg(feature = "threading")] fn remove(&mut self, request: &SectorGroupRequest) { self.inflight.remove(request); } + #[cfg(feature = "threading")] fn contains(&self, request: &SectorGroupRequest) -> bool { self.lru_cache.contains(request) || self.inflight.contains_key(request) } } #[allow(unused)] +#[cfg(feature = "threading")] struct PreloaderThreadStats { thread_id: usize, wait_time: Duration, @@ -187,6 +206,7 @@ struct PreloaderThreadStats { io_time: Duration, } +#[cfg(feature = "threading")] fn preloader_thread( thread_id: usize, request_rx: Receiver, @@ -237,6 +257,7 @@ fn preloader_thread( } impl Preloader { + #[cfg(feature = "threading")] pub fn new(loader: SectorGroupLoader, num_threads: usize) -> Arc { debug!("Creating preloader with {} threads", num_threads); @@ -258,56 +279,73 @@ impl Preloader { Arc::new(Self { request_tx, request_rx, stat_tx, stat_rx, threads, cache, loader }) } + #[cfg(not(feature = "threading"))] + pub fn new(loader: SectorGroupLoader) -> Arc { + debug!("Creating single-threaded preloader"); + let cache = Arc::new(Mutex::new(PreloaderCache::default())); + let loader = Mutex::new(loader); + Arc::new(Self { cache, loader }) + } + #[allow(unused)] pub fn shutdown(self) { - let guard = self.threads.into_inner().unwrap(); - for handle in guard.join_handles { - handle.join().unwrap(); + #[cfg(feature = "threading")] + { + let guard = self.threads.into_inner().unwrap(); + for handle in guard.join_handles { + handle.join().unwrap(); + } } } #[instrument(name = "Preloader::fetch", skip_all)] pub fn fetch(&self, request: SectorGroupRequest, max_groups: u32) -> SectorGroupResult { - let num_threads = { - let mut threads_guard = self.threads.lock().map_err(map_poisoned)?; - while let Ok(stat) = self.stat_rx.try_recv() { - threads_guard.push_stats(stat, self); - } - threads_guard.join_handles.len() - }; - let mut cache_guard = self.cache.lock().map_err(map_poisoned)?; - // Preload n groups ahead - for i in 0..num_threads as u32 { - let group_idx = request.group_idx + i; - if group_idx >= max_groups { - break; - } - let request = SectorGroupRequest { group_idx, ..request }; - if cache_guard.contains(&request) { - continue; - } - if self.request_tx.send(request).is_ok() { - cache_guard.inflight.insert(request, WaitGroup::new()); - } - } - if let Some(cached) = cache_guard.lru_cache.get(&request) { - return Ok(cached.clone()); - } - if let Some(wg) = cache_guard.inflight.get(&request) { - // Wait for inflight request to finish - let wg = wg.clone(); - drop(cache_guard); - { - let _span = span!(Level::TRACE, "wg.wait").entered(); - wg.wait(); - } + #[cfg(feature = "threading")] + { + let num_threads = { + let mut threads_guard = self.threads.lock().map_err(map_poisoned)?; + while let Ok(stat) = self.stat_rx.try_recv() { + threads_guard.push_stats(stat, self); + } + threads_guard.join_handles.len() + }; let mut cache_guard = self.cache.lock().map_err(map_poisoned)?; + // Preload n groups ahead + for i in 0..num_threads as u32 { + let group_idx = request.group_idx + i; + if group_idx >= max_groups { + break; + } + let request = SectorGroupRequest { group_idx, ..request }; + if cache_guard.contains(&request) { + continue; + } + if self.request_tx.send(request).is_ok() { + cache_guard.inflight.insert(request, WaitGroup::new()); + } + } if let Some(cached) = cache_guard.lru_cache.get(&request) { return Ok(cached.clone()); } - } else { - drop(cache_guard); + if let Some(wg) = cache_guard.inflight.get(&request) { + // Wait for inflight request to finish + let wg = wg.clone(); + drop(cache_guard); + { + let _span = span!(Level::TRACE, "wg.wait").entered(); + wg.wait(); + } + let mut cache_guard = self.cache.lock().map_err(map_poisoned)?; + if let Some(cached) = cache_guard.lru_cache.get(&request) { + return Ok(cached.clone()); + } + } else { + drop(cache_guard); + } } + #[cfg(not(feature = "threading"))] + let _ = max_groups; + // No threads are running, fallback to single-threaded loader let result = { let mut loader = self.loader.lock().map_err(map_poisoned)?; diff --git a/nod/src/disc/reader.rs b/nod/src/disc/reader.rs index 6a94b38..2da1b79 100644 --- a/nod/src/disc/reader.rs +++ b/nod/src/disc/reader.rs @@ -147,6 +147,7 @@ impl DiscReader { let size = io.meta().disc_size.unwrap_or_else(|| guess_disc_size(partitions)); let preloader = Preloader::new( SectorGroupLoader::new(io.clone(), disc_header_arc, partitions.clone()), + #[cfg(feature = "threading")] options.preloader_threads, ); Ok(Self { diff --git a/nod/src/disc/writer.rs b/nod/src/disc/writer.rs index 30a8ef0..d758c4c 100644 --- a/nod/src/disc/writer.rs +++ b/nod/src/disc/writer.rs @@ -1,5 +1,4 @@ use std::{ - collections::VecDeque, io, io::{BufRead, Read}, }; @@ -8,7 +7,7 @@ use bytes::{Bytes, BytesMut}; use dyn_clone::DynClone; use crate::{ - Error, Result, ResultContext, + Result, ResultContext, common::{PartitionInfo, PartitionKind}, disc::{ SECTOR_SIZE, @@ -97,41 +96,51 @@ pub fn read_block(reader: &mut DiscReader, block_size: usize) -> io::Result<(Byt } /// Process blocks in parallel, ensuring that they are written in order. +#[cfg_attr(not(feature = "threading"), inline)] pub(crate) fn par_process( mut processor: P, block_count: u32, - num_threads: usize, + #[cfg(feature = "threading")] num_threads: usize, mut callback: impl FnMut(BlockResult) -> Result<()>, ) -> Result<()> where T: Send, P: BlockProcessor, { - if num_threads == 0 { - // Fall back to single-threaded processing - for block_idx in 0..block_count { - let block = processor - .process_block(block_idx) - .with_context(|| format!("Failed to process block {block_idx}"))?; - callback(block)?; - } - return Ok(()); - } + #[cfg(feature = "threading")] + if num_threads > 0 { + return std::thread::scope(|s| { + use std::collections::VecDeque; - std::thread::scope(|s| { - let (block_tx, block_rx) = crossbeam_channel::bounded(block_count as usize); - for block_idx in 0..block_count { - block_tx.send(block_idx).unwrap(); - } - drop(block_tx); // Disconnect channel + use crate::Error; - let (result_tx, result_rx) = crossbeam_channel::bounded(0); + let (block_tx, block_rx) = crossbeam_channel::bounded(block_count as usize); + for block_idx in 0..block_count { + block_tx.send(block_idx).unwrap(); + } + drop(block_tx); // Disconnect channel - // Spawn threads to process blocks - for _ in 0..num_threads - 1 { - let block_rx = block_rx.clone(); - let result_tx = result_tx.clone(); - let mut processor = processor.clone(); + let (result_tx, result_rx) = crossbeam_channel::bounded(0); + + // Spawn threads to process blocks + for _ in 0..num_threads - 1 { + let block_rx = block_rx.clone(); + let result_tx = result_tx.clone(); + let mut processor = processor.clone(); + s.spawn(move || { + while let Ok(block_idx) = block_rx.recv() { + let result = processor + .process_block(block_idx) + .with_context(|| format!("Failed to process block {block_idx}")); + let failed = result.is_err(); // Stop processing if an error occurs + if result_tx.send(result).is_err() || failed { + break; + } + } + }); + } + + // Last iteration moves instead of cloning s.spawn(move || { while let Ok(block_idx) = block_rx.recv() { let result = processor @@ -143,45 +152,41 @@ where } } }); - } - // Last iteration moves instead of cloning - s.spawn(move || { - while let Ok(block_idx) = block_rx.recv() { - let result = processor - .process_block(block_idx) - .with_context(|| format!("Failed to process block {block_idx}")); - let failed = result.is_err(); // Stop processing if an error occurs - if result_tx.send(result).is_err() || failed { - break; - } - } - }); - - // Main thread processes results - let mut current_block = 0; - let mut out_of_order = VecDeque::>::new(); - while let Ok(result) = result_rx.recv() { - let result = result?; - if result.block_idx == current_block { - callback(result)?; - current_block += 1; - // Check if any out of order blocks can be written - while out_of_order.front().is_some_and(|r| r.block_idx == current_block) { - callback(out_of_order.pop_front().unwrap())?; + // Main thread processes results + let mut current_block = 0; + let mut out_of_order = VecDeque::>::new(); + while let Ok(result) = result_rx.recv() { + let result = result?; + if result.block_idx == current_block { + callback(result)?; current_block += 1; - } - } else { - // Insert sorted - match out_of_order.binary_search_by_key(&result.block_idx, |r| r.block_idx) { - Ok(idx) => Err(Error::Other(format!("Unexpected duplicate block {idx}")))?, - Err(idx) => out_of_order.insert(idx, result), + // Check if any out of order blocks can be written + while out_of_order.front().is_some_and(|r| r.block_idx == current_block) { + callback(out_of_order.pop_front().unwrap())?; + current_block += 1; + } + } else { + // Insert sorted + match out_of_order.binary_search_by_key(&result.block_idx, |r| r.block_idx) { + Ok(idx) => Err(Error::Other(format!("Unexpected duplicate block {idx}")))?, + Err(idx) => out_of_order.insert(idx, result), + } } } - } - Ok(()) - }) + Ok(()) + }); + } + + // Fall back to single-threaded processing + for block_idx in 0..block_count { + let block = processor + .process_block(block_idx) + .with_context(|| format!("Failed to process block {block_idx}"))?; + callback(block)?; + } + return Ok(()); } /// The determined block type. diff --git a/nod/src/io/ciso.rs b/nod/src/io/ciso.rs index 67eddb3..42ff99e 100644 --- a/nod/src/io/ciso.rs +++ b/nod/src/io/ciso.rs @@ -272,6 +272,7 @@ impl DiscWriter for DiscWriterCISO { scrub_update_partition: options.scrub == ScrubLevel::UpdatePartition, }, self.block_count, + #[cfg(feature = "threading")] options.processor_threads, |block| -> Result<()> { // Update hashers diff --git a/nod/src/io/gcz.rs b/nod/src/io/gcz.rs index 7325363..eb10ce7 100644 --- a/nod/src/io/gcz.rs +++ b/nod/src/io/gcz.rs @@ -325,6 +325,7 @@ impl DiscWriter for DiscWriterGCZ { compressor: Compressor::new(self.compression, block_size as usize), }, block_count, + #[cfg(feature = "threading")] options.processor_threads, |block| { // Update hashers diff --git a/nod/src/io/wbfs.rs b/nod/src/io/wbfs.rs index 1dc42ba..600dc47 100644 --- a/nod/src/io/wbfs.rs +++ b/nod/src/io/wbfs.rs @@ -325,6 +325,7 @@ impl DiscWriter for DiscWriterWBFS { scrub_update_partition: options.scrub == ScrubLevel::UpdatePartition, }, self.block_count as u32, + #[cfg(feature = "threading")] options.processor_threads, |block| -> Result<()> { // Update hashers diff --git a/nod/src/io/wia.rs b/nod/src/io/wia.rs index 33575f3..8adaed5 100644 --- a/nod/src/io/wia.rs +++ b/nod/src/io/wia.rs @@ -1729,6 +1729,7 @@ impl DiscWriter for DiscWriterWIA { junk_info: self.junk_info.clone(), }, self.group_count, + #[cfg(feature = "threading")] options.processor_threads, |group| -> Result<()> { // Update hashers diff --git a/nod/src/read.rs b/nod/src/read.rs index 9531d06..7eff011 100644 --- a/nod/src/read.rs +++ b/nod/src/read.rs @@ -52,6 +52,7 @@ pub struct DiscOptions { /// is particularly useful when reading the disc image sequentially, as it /// can perform decompression and rebuilding in parallel with the main /// read thread. The default value of 0 disables preloading. + #[cfg(feature = "threading")] pub preloader_threads: usize, } diff --git a/nod/src/util/digest.rs b/nod/src/util/digest.rs index a1191c4..2b16fea 100644 --- a/nod/src/util/digest.rs +++ b/nod/src/util/digest.rs @@ -1,7 +1,4 @@ -use std::{thread, thread::JoinHandle}; - use bytes::Bytes; -use crossbeam_channel::Sender; use digest::Digest; use tracing::instrument; @@ -33,85 +30,156 @@ pub fn sha1_hash(buf: &[u8]) -> HashBytes { #[instrument(skip_all)] pub fn xxh64_hash(buf: &[u8]) -> u64 { xxhash_rust::xxh64::xxh64(buf, 0) } -pub type DigestThread = (Sender, JoinHandle); +#[cfg(feature = "threading")] +mod multi_threaded { + use std::{thread, thread::JoinHandle}; -pub fn digest_thread() -> DigestThread -where H: Hasher + Send + 'static { - let (tx, rx) = crossbeam_channel::bounded::(1); - let handle = thread::Builder::new() - .name(format!("Digest {}", H::NAME)) - .spawn(move || { - let mut hasher = H::new(); - while let Ok(data) = rx.recv() { - hasher.update(data.as_ref()); - } - hasher.finalize() - }) - .expect("Failed to spawn digest thread"); - (tx, handle) -} + use crossbeam_channel::Sender; -pub struct DigestManager { - threads: Vec, -} + use super::*; -impl DigestManager { - pub fn new(options: &ProcessOptions) -> Self { - let mut threads = Vec::new(); - if options.digest_crc32 { - threads.push(digest_thread::()); - } - if options.digest_md5 { - #[cfg(feature = "openssl")] - threads.push(digest_thread::()); - #[cfg(not(feature = "openssl"))] - threads.push(digest_thread::()); - } - if options.digest_sha1 { - #[cfg(feature = "openssl")] - threads.push(digest_thread::()); - #[cfg(not(feature = "openssl"))] - threads.push(digest_thread::()); - } - if options.digest_xxh64 { - threads.push(digest_thread::()); - } - DigestManager { threads } + type DigestThread = (Sender, JoinHandle); + + fn digest_thread() -> DigestThread + where H: Hasher + Send + 'static { + let (tx, rx) = crossbeam_channel::bounded::(1); + let handle = thread::Builder::new() + .name(format!("Digest {}", H::NAME)) + .spawn(move || { + let mut hasher = H::new(); + while let Ok(data) = rx.recv() { + hasher.update(data.as_ref()); + } + hasher.finalize() + }) + .expect("Failed to spawn digest thread"); + (tx, handle) } - #[instrument(name = "DigestManager::send", skip_all)] - pub fn send(&self, data: Bytes) { - let mut sent = 0usize; - // Non-blocking send to all threads - for (idx, (tx, _)) in self.threads.iter().enumerate() { - if tx.try_send(data.clone()).is_ok() { - sent |= 1 << idx; - } - } - // Blocking send to any remaining threads - for (idx, (tx, _)) in self.threads.iter().enumerate() { - if sent & (1 << idx) == 0 { - tx.send(data.clone()).expect("Failed to send data to digest thread"); - } - } + pub struct DigestManager { + threads: Vec, } - #[instrument(name = "DigestManager::finish", skip_all)] - pub fn finish(self) -> DigestResults { - let mut results = DigestResults { crc32: None, md5: None, sha1: None, xxh64: None }; - for (tx, handle) in self.threads { - drop(tx); // Close channel - match handle.join().unwrap() { - DigestResult::Crc32(v) => results.crc32 = Some(v), - DigestResult::Md5(v) => results.md5 = Some(v), - DigestResult::Sha1(v) => results.sha1 = Some(v), - DigestResult::Xxh64(v) => results.xxh64 = Some(v), + impl DigestManager { + pub fn new(options: &ProcessOptions) -> Self { + let mut threads = Vec::new(); + if options.digest_crc32 { + threads.push(digest_thread::()); + } + if options.digest_md5 { + #[cfg(feature = "openssl")] + threads.push(digest_thread::()); + #[cfg(not(feature = "openssl"))] + threads.push(digest_thread::()); + } + if options.digest_sha1 { + #[cfg(feature = "openssl")] + threads.push(digest_thread::()); + #[cfg(not(feature = "openssl"))] + threads.push(digest_thread::()); + } + if options.digest_xxh64 { + threads.push(digest_thread::()); + } + DigestManager { threads } + } + + #[instrument(name = "DigestManager::send", skip_all)] + pub fn send(&self, data: Bytes) { + let mut sent = 0usize; + // Non-blocking send to all threads + for (idx, (tx, _)) in self.threads.iter().enumerate() { + if tx.try_send(data.clone()).is_ok() { + sent |= 1 << idx; + } + } + // Blocking send to any remaining threads + for (idx, (tx, _)) in self.threads.iter().enumerate() { + if sent & (1 << idx) == 0 { + tx.send(data.clone()).expect("Failed to send data to digest thread"); + } } } - results + + #[instrument(name = "DigestManager::finish", skip_all)] + pub fn finish(self) -> DigestResults { + let mut results = DigestResults { crc32: None, md5: None, sha1: None, xxh64: None }; + for (tx, handle) in self.threads { + drop(tx); // Close channel + match handle.join().unwrap() { + DigestResult::Crc32(v) => results.crc32 = Some(v), + DigestResult::Md5(v) => results.md5 = Some(v), + DigestResult::Sha1(v) => results.sha1 = Some(v), + DigestResult::Xxh64(v) => results.xxh64 = Some(v), + } + } + results + } } } +#[cfg(not(feature = "threading"))] +mod single_threaded { + use std::cell::RefCell; + + use super::*; + + pub struct DigestManager { + hashers: Vec>>, + } + + impl DigestManager { + pub fn new(options: &ProcessOptions) -> Self { + let mut hashers = Vec::>>::new(); + if options.digest_crc32 { + hashers.push(RefCell::new(Box::new(crc32fast::Hasher::new()))); + } + if options.digest_md5 { + #[cfg(feature = "openssl")] + hashers.push(RefCell::new(Box::new(openssl_util::HasherMD5::new()))); + #[cfg(not(feature = "openssl"))] + hashers.push(RefCell::new(Box::new(md5::Md5::new()))); + } + if options.digest_sha1 { + #[cfg(feature = "openssl")] + hashers.push(RefCell::new(Box::new(openssl_util::HasherSHA1::new()))); + #[cfg(not(feature = "openssl"))] + hashers.push(RefCell::new(Box::new(sha1::Sha1::new()))); + } + if options.digest_xxh64 { + hashers.push(RefCell::new(Box::new(xxhash_rust::xxh64::Xxh64::new(0)))); + } + Self { hashers } + } + + #[instrument(name = "DigestManager::send", skip_all)] + pub fn send(&self, data: Bytes) { + for hasher in &self.hashers { + hasher.borrow_mut().update(&data); + } + } + + #[instrument(name = "DigestManager::finish", skip_all)] + pub fn finish(self) -> DigestResults { + let mut results = DigestResults { crc32: None, md5: None, sha1: None, xxh64: None }; + for hasher in self.hashers { + match hasher.borrow_mut().finalize() { + DigestResult::Crc32(v) => results.crc32 = Some(v), + DigestResult::Md5(v) => results.md5 = Some(v), + DigestResult::Sha1(v) => results.sha1 = Some(v), + DigestResult::Xxh64(v) => results.xxh64 = Some(v), + } + } + results + } + } +} + +#[cfg(feature = "threading")] +pub use multi_threaded::DigestManager; +#[cfg(not(feature = "threading"))] +pub use single_threaded::DigestManager; + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum DigestResult { Crc32(u32), @@ -121,19 +189,25 @@ pub enum DigestResult { } pub trait Hasher { + #[cfg(feature = "threading")] const NAME: &'static str; + #[cfg(feature = "threading")] fn new() -> Self; - fn finalize(self) -> DigestResult; + fn finalize(&mut self) -> DigestResult; fn update(&mut self, data: &[u8]); } impl Hasher for md5::Md5 { + #[cfg(feature = "threading")] const NAME: &'static str = "MD5"; + #[cfg(feature = "threading")] fn new() -> Self { Digest::new() } - fn finalize(self) -> DigestResult { DigestResult::Md5(Digest::finalize(self).into()) } + fn finalize(&mut self) -> DigestResult { + DigestResult::Md5(Digest::finalize_reset(self).into()) + } #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347 #[instrument(name = "md5::Md5::update", skip_all)] @@ -141,11 +215,15 @@ impl Hasher for md5::Md5 { } impl Hasher for sha1::Sha1 { + #[cfg(feature = "threading")] const NAME: &'static str = "SHA-1"; + #[cfg(feature = "threading")] fn new() -> Self { Digest::new() } - fn finalize(self) -> DigestResult { DigestResult::Sha1(Digest::finalize(self).into()) } + fn finalize(&mut self) -> DigestResult { + DigestResult::Sha1(Digest::finalize_reset(self).into()) + } #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347 #[instrument(name = "sha1::Sha1::update", skip_all)] @@ -153,11 +231,15 @@ impl Hasher for sha1::Sha1 { } impl Hasher for crc32fast::Hasher { + #[cfg(feature = "threading")] const NAME: &'static str = "CRC32"; + #[cfg(feature = "threading")] fn new() -> Self { crc32fast::Hasher::new() } - fn finalize(self) -> DigestResult { DigestResult::Crc32(crc32fast::Hasher::finalize(self)) } + fn finalize(&mut self) -> DigestResult { + DigestResult::Crc32(crc32fast::Hasher::finalize(self.clone())) + } #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347 #[instrument(name = "crc32fast::Hasher::update", skip_all)] @@ -165,12 +247,14 @@ impl Hasher for crc32fast::Hasher { } impl Hasher for xxhash_rust::xxh64::Xxh64 { + #[cfg(feature = "threading")] const NAME: &'static str = "XXH64"; + #[cfg(feature = "threading")] fn new() -> Self { xxhash_rust::xxh64::Xxh64::new(0) } - fn finalize(self) -> DigestResult { - DigestResult::Xxh64(xxhash_rust::xxh64::Xxh64::digest(&self)) + fn finalize(&mut self) -> DigestResult { + DigestResult::Xxh64(xxhash_rust::xxh64::Xxh64::digest(self)) } #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347 @@ -197,7 +281,7 @@ mod openssl_util { impl HashWrapper where T: MessageDigest { - fn new() -> Self { + pub(super) fn new() -> Self { Self { hasher: openssl::hash::Hasher::new(T::new()).unwrap(), _marker: Default::default(), @@ -222,11 +306,13 @@ mod openssl_util { } impl Hasher for HasherMD5 { + #[cfg(feature = "threading")] const NAME: &'static str = "MD5"; + #[cfg(feature = "threading")] fn new() -> Self { Self::new() } - fn finalize(mut self) -> DigestResult { + fn finalize(&mut self) -> DigestResult { DigestResult::Md5((*self.hasher.finish().unwrap()).try_into().unwrap()) } @@ -236,11 +322,13 @@ mod openssl_util { } impl Hasher for HasherSHA1 { + #[cfg(feature = "threading")] const NAME: &'static str = "SHA-1"; + #[cfg(feature = "threading")] fn new() -> Self { Self::new() } - fn finalize(mut self) -> DigestResult { + fn finalize(&mut self) -> DigestResult { DigestResult::Sha1((*self.hasher.finish().unwrap()).try_into().unwrap()) } diff --git a/nod/src/write.rs b/nod/src/write.rs index 11e8cba..01b17ed 100644 --- a/nod/src/write.rs +++ b/nod/src/write.rs @@ -45,30 +45,31 @@ pub struct ProcessOptions { /// If the output format supports multithreaded processing, this sets the number of threads to /// use for processing data. This is particularly useful for formats that compress data or /// perform other transformations. The default value of 0 disables multithreading. + #[cfg(feature = "threading")] pub processor_threads: usize, /// Enables CRC32 checksum calculation for the disc data. /// /// If the output format supports it, this will be stored in the disc data. (NKit 2 compatible) - /// Each digest calculation will run on a separate thread, unaffected by the processor thread - /// count. + /// If the "threading" feature is enabled, each digest calculation will run on a separate thread, + /// unaffected by the processor thread count. pub digest_crc32: bool, /// Enables MD5 checksum calculation for the disc data. (Slow!) /// /// If the output format supports it, this will be stored in the disc data. (NKit 2 compatible) - /// Each digest calculation will run on a separate thread, unaffected by the processor thread - /// count. + /// If the "threading" feature is enabled, each digest calculation will run on a separate thread, + /// unaffected by the processor thread count. pub digest_md5: bool, /// Enables SHA-1 checksum calculation for the disc data. /// /// If the output format supports it, this will be stored in the disc data. (NKit 2 compatible) - /// Each digest calculation will run on a separate thread, unaffected by the processor thread - /// count. + /// If the "threading" feature is enabled, each digest calculation will run on a separate thread, + /// unaffected by the processor thread count. pub digest_sha1: bool, /// Enables XXH64 checksum calculation for the disc data. /// /// If the output format supports it, this will be stored in the disc data. (NKit 2 compatible) - /// Each digest calculation will run on a separate thread, unaffected by the processor thread - /// count. + /// If the "threading" feature is enabled, each digest calculation will run on a separate thread, + /// unaffected by the processor thread count. pub digest_xxh64: bool, /// The level of scrubbing to perform on the disc image. /// diff --git a/nodtool/Cargo.toml b/nodtool/Cargo.toml index 6d9b1d1..4a6ad55 100644 --- a/nodtool/Cargo.toml +++ b/nodtool/Cargo.toml @@ -16,7 +16,7 @@ categories = ["command-line-utilities", "parser-implementations"] build = "build.rs" [features] -default = ["compress-bzip2", "compress-lzma", "compress-zlib", "compress-zstd"] +default = ["compress-bzip2", "compress-lzma", "compress-zlib", "compress-zstd", "threading"] compress-bzip2 = ["nod/compress-bzip2"] compress-lzma = ["nod/compress-lzma"] compress-zlib = ["nod/compress-zlib"] @@ -24,6 +24,7 @@ compress-zstd = ["nod/compress-zstd"] openssl = ["nod/openssl"] openssl-vendored = ["nod/openssl-vendored"] tracy = ["dep:tracing-tracy"] +threading = ["nod/threading", "dep:num_cpus"] [dependencies] argp = "0.4" @@ -34,7 +35,7 @@ hex = { version = "0.4", features = ["serde"] } indicatif = "0.18" md-5 = { workspace = true } nod = { version = "2.0.0-alpha", path = "../nod", default-features = false } -num_cpus = "1.17" +num_cpus = { version = "1.17", optional = true } quick-xml = { version = "0.38", features = ["serialize"] } serde = { version = "1.0", features = ["derive"] } sha1 = { workspace = true } diff --git a/nodtool/src/cmd/convert.rs b/nodtool/src/cmd/convert.rs index 065ef34..5d5436d 100644 --- a/nodtool/src/cmd/convert.rs +++ b/nodtool/src/cmd/convert.rs @@ -52,6 +52,7 @@ pub fn run(args: Args) -> nod::Result<()> { )); } }, + #[cfg(feature = "threading")] preloader_threads: 4, }; let format = match args.out.extension() { diff --git a/nodtool/src/cmd/dat.rs b/nodtool/src/cmd/dat.rs index 9e81790..038b149 100644 --- a/nodtool/src/cmd/dat.rs +++ b/nodtool/src/cmd/dat.rs @@ -160,8 +160,11 @@ struct DiscHashes { } fn load_disc(path: &Path, name: &str, full_verify: bool) -> Result { - let options = - DiscOptions { partition_encryption: PartitionEncryption::Original, preloader_threads: 4 }; + let options = DiscOptions { + partition_encryption: PartitionEncryption::Original, + #[cfg(feature = "threading")] + preloader_threads: 4, + }; let disc = DiscReader::new(path, &options)?; if !full_verify { let meta = disc.meta(); @@ -187,6 +190,7 @@ fn load_disc(path: &Path, name: &str, full_verify: bool) -> Result { Ok(()) }, &ProcessOptions { + #[cfg(feature = "threading")] processor_threads: 12, // TODO digest_crc32: true, digest_md5: false, diff --git a/nodtool/src/cmd/extract.rs b/nodtool/src/cmd/extract.rs index 3c8eacf..88a5496 100644 --- a/nodtool/src/cmd/extract.rs +++ b/nodtool/src/cmd/extract.rs @@ -53,8 +53,11 @@ pub fn run(args: Args) -> nod::Result<()> { } else { output_dir = args.file.with_extension(""); } - let disc = - DiscReader::new(&args.file, &DiscOptions { preloader_threads: 4, ..Default::default() })?; + let disc = DiscReader::new(&args.file, &DiscOptions { + #[cfg(feature = "threading")] + preloader_threads: 4, + ..Default::default() + })?; let header = disc.header(); let is_wii = header.is_wii(); let options = PartitionOptions { validate_hashes: args.validate }; diff --git a/nodtool/src/cmd/verify.rs b/nodtool/src/cmd/verify.rs index fb23480..fa04d3d 100644 --- a/nodtool/src/cmd/verify.rs +++ b/nodtool/src/cmd/verify.rs @@ -34,7 +34,6 @@ pub fn run(args: Args) -> nod::Result<()> { println!("Loading dat files..."); redump::load_dats(args.dat.iter().map(PathBuf::as_ref))?; } - let cpus = num_cpus::get(); let options = DiscOptions { partition_encryption: match (args.decrypt, args.encrypt) { (true, false) => PartitionEncryption::ForceDecrypted, @@ -46,7 +45,8 @@ pub fn run(args: Args) -> nod::Result<()> { )); } }, - preloader_threads: 4.min(cpus), + #[cfg(feature = "threading")] + preloader_threads: 4.min(num_cpus::get()), }; let format_options = FormatOptions::default(); for file in &args.file { diff --git a/nodtool/src/util/shared.rs b/nodtool/src/util/shared.rs index 64348d1..ff0210c 100644 --- a/nodtool/src/util/shared.rs +++ b/nodtool/src/util/shared.rs @@ -11,7 +11,7 @@ use nod::{ common::Compression, disc::DiscHeader, read::{DiscMeta, DiscOptions, DiscReader, PartitionEncryption}, - write::{DiscWriter, DiscWriterWeight, FormatOptions, ProcessOptions, ScrubLevel}, + write::{DiscWriter, FormatOptions, ProcessOptions, ScrubLevel}, }; use size::Size; @@ -100,11 +100,15 @@ pub fn convert_and_verify( }) .progress_chars("#>-")); - let cpus = num_cpus::get(); - let processor_threads = match disc_writer.weight() { - DiscWriterWeight::Light => 0, - DiscWriterWeight::Medium => cpus / 2, - DiscWriterWeight::Heavy => cpus, + #[cfg(feature = "threading")] + let processor_threads = { + use nod::write::DiscWriterWeight; + let cpus = num_cpus::get(); + match disc_writer.weight() { + DiscWriterWeight::Light => 0, + DiscWriterWeight::Medium => cpus / 2, + DiscWriterWeight::Heavy => cpus, + } }; let mut total_written = 0u64; @@ -118,6 +122,7 @@ pub fn convert_and_verify( Ok(()) }, &ProcessOptions { + #[cfg(feature = "threading")] processor_threads, digest_crc32: true, digest_md5: md5,