encounter/nod
0
0
Fork 0
mirror of https://github.com/encounter/nod.git synced 2026-07-10 12:18:48 -07:00
Code Issues Packages Projects Releases Wiki Activity

Support single-threaded runtimes with optional "threading" feature

Browse Source
This commit is contained in:
Luke Street
2025-12-29 16:07:28 -07:00
parent 6fff7cb818
commit d2c7e0ec40
17 changed files with 358 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nod/Cargo.toml
+4 -3
View File
@@ -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"
nod/src/disc/preloader.rs
+80 -42
View File
@@ -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<Duration>,
#[allow(unused)] // TODO WIA hash exceptions
pub group_hashes: Option<Arc<GroupHashes>>,
@@ -76,10 +82,15 @@ pub type SectorGroupResult = io::Result<SectorGroup>;
#[allow(unused)]
pub struct Preloader {
#[cfg(feature = "threading")]
request_tx: Sender<SectorGroupRequest>,
#[cfg(feature = "threading")]
request_rx: Receiver<SectorGroupRequest>,
#[cfg(feature = "threading")]
stat_tx: Sender<PreloaderThreadStats>,
#[cfg(feature = "threading")]
stat_rx: Receiver<PreloaderThreadStats>,
#[cfg(feature = "threading")]
threads: Mutex<PreloaderThreads>,
cache: Arc<Mutex<PreloaderCache>>,
// Fallback single-threaded loader
@@ -87,6 +98,7 @@ pub struct Preloader {
}
#[allow(unused)]
#[cfg(feature = "threading")]
struct PreloaderThreads {
join_handles: Vec<JoinHandle<()>>,
last_adjust: Instant,
@@ -96,6 +108,7 @@ struct PreloaderThreads {
io_time_avg: SingleSumSMA<Duration, u32, 100>,
}
#[cfg(feature = "threading")]
impl PreloaderThreads {
fn new(join_handles: Vec<JoinHandle<()>>) -> Self {
Self {
@@ -153,6 +166,7 @@ impl PreloaderThreads {
}
struct PreloaderCache {
#[cfg(feature = "threading")]
inflight: HashMap<SectorGroupRequest, WaitGroup>,
lru_cache: LruCache<SectorGroupRequest, SectorGroup>,
}
@@ -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<SectorGroupRequest>,
@@ -237,6 +257,7 @@ fn preloader_thread(
}
impl Preloader {
#[cfg(feature = "threading")]
pub fn new(loader: SectorGroupLoader, num_threads: usize) -> Arc<Self> {
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<Self> {
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)?;
nod/src/disc/reader.rs
+1
View File
@@ -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 {
nod/src/disc/writer.rs
+64 -59
View File
@@ -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<P, T>(
mut processor: P,
block_count: u32,
num_threads: usize,
#[cfg(feature = "threading")] num_threads: usize,
mut callback: impl FnMut(BlockResult<T>) -> Result<()>,
) -> Result<()>
where
T: Send,
P: BlockProcessor<BlockMeta = T>,
{
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::<BlockResult<T>>::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::<BlockResult<T>>::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.
nod/src/io/ciso.rs
+1
View File
@@ -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
nod/src/io/gcz.rs
+1
View File
@@ -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
nod/src/io/wbfs.rs
+1
View File
@@ -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
nod/src/io/wia.rs
+1
View File
@@ -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
nod/src/read.rs
+1
View File
@@ -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,
}
nod/src/util/digest.rs
+167 -79
View File
@@ -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<Bytes>, JoinHandle<DigestResult>);
#[cfg(feature = "threading")]
mod multi_threaded {
use std::{thread, thread::JoinHandle};
pub fn digest_thread<H>() -> DigestThread
where H: Hasher + Send + 'static {
let (tx, rx) = crossbeam_channel::bounded::<Bytes>(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<DigestThread>,
}
use super::*;
impl DigestManager {
pub fn new(options: &ProcessOptions) -> Self {
let mut threads = Vec::new();
if options.digest_crc32 {
threads.push(digest_thread::<crc32fast::Hasher>());
}
if options.digest_md5 {
#[cfg(feature = "openssl")]
threads.push(digest_thread::<openssl_util::HasherMD5>());
#[cfg(not(feature = "openssl"))]
threads.push(digest_thread::<md5::Md5>());
}
if options.digest_sha1 {
#[cfg(feature = "openssl")]
threads.push(digest_thread::<openssl_util::HasherSHA1>());
#[cfg(not(feature = "openssl"))]
threads.push(digest_thread::<sha1::Sha1>());
}
if options.digest_xxh64 {
threads.push(digest_thread::<xxhash_rust::xxh64::Xxh64>());
}
DigestManager { threads }
type DigestThread = (Sender<Bytes>, JoinHandle<DigestResult>);
fn digest_thread<H>() -> DigestThread
where H: Hasher + Send + 'static {
let (tx, rx) = crossbeam_channel::bounded::<Bytes>(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<DigestThread>,
}
#[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::<crc32fast::Hasher>());
}
if options.digest_md5 {
#[cfg(feature = "openssl")]
threads.push(digest_thread::<openssl_util::HasherMD5>());
#[cfg(not(feature = "openssl"))]
threads.push(digest_thread::<md5::Md5>());
}
if options.digest_sha1 {
#[cfg(feature = "openssl")]
threads.push(digest_thread::<openssl_util::HasherSHA1>());
#[cfg(not(feature = "openssl"))]
threads.push(digest_thread::<sha1::Sha1>());
}
if options.digest_xxh64 {
threads.push(digest_thread::<xxhash_rust::xxh64::Xxh64>());
}
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<RefCell<Box<dyn Hasher>>>,
}
impl DigestManager {
pub fn new(options: &ProcessOptions) -> Self {
let mut hashers = Vec::<RefCell<Box<dyn Hasher>>>::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<T> HashWrapper<T>
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())
}
nod/src/write.rs
+9 -8
View File
@@ -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.
///
nodtool/Cargo.toml
+3 -2
View File
@@ -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 }
nodtool/src/cmd/convert.rs
+1
View File
@@ -52,6 +52,7 @@ pub fn run(args: Args) -> nod::Result<()> {
));
}
},
#[cfg(feature = "threading")]
preloader_threads: 4,
};
let format = match args.out.extension() {
nodtool/src/cmd/dat.rs
+6 -2
View File
@@ -160,8 +160,11 @@ struct DiscHashes {
}
fn load_disc(path: &Path, name: &str, full_verify: bool) -> Result<DiscHashes> {
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<DiscHashes> {
Ok(())
},
&ProcessOptions {
#[cfg(feature = "threading")]
processor_threads: 12, // TODO
digest_crc32: true,
digest_md5: false,
nodtool/src/cmd/extract.rs
+5 -2
View File
@@ -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 };
nodtool/src/cmd/verify.rs
+2 -2
View File
@@ -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 {
nodtool/src/util/shared.rs
+11 -6
View File
@@ -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,