#![forbid(unsafe_code)]
mod args;
use anyhow as ah;
use args::{Args, parse_args};
use chrono::prelude::*;
use disktest_lib::{Disktest, DisktestFile, DisktestQuiet};
use std::{
env::args_os,
sync::{Arc, atomic::AtomicBool},
};
fn install_abort_handlers() -> ah::Result<Arc<AtomicBool>> {
let abort = Arc::new(AtomicBool::new(false));
for sig in &[
signal_hook::consts::signal::SIGTERM,
signal_hook::consts::signal::SIGINT,
] {
if let Err(e) = signal_hook::flag::register(*sig, Arc::clone(&abort)) {
return Err(ah::format_err!("Failed to register signal {sig}: {e}"));
}
}
Ok(abort)
}
fn print_generated_seed(seed: &str, verbose: bool) {
if verbose {
println!(
"\nThe generated --seed is:\n {seed}\nUse this seed for subsequent --verify.\n"
);
} else {
println!("Generated --seed {seed}\n");
}
}
fn new_disktest(
args: &Args,
round_id: u64,
write: bool,
abort: &Arc<AtomicBool>,
) -> ah::Result<(Disktest, DisktestFile)> {
Ok((
Disktest::new(
args.algorithm,
args.seed.as_bytes(),
round_id,
args.invert_pattern,
args.threads,
args.quiet,
Some(Arc::clone(abort)),
),
DisktestFile::open(&args.device, !write, write)?,
))
}
fn main() -> ah::Result<()> {
let args = parse_args(args_os())?;
let abort = install_abort_handlers()?;
if !args.user_seed && args.quiet < DisktestQuiet::NoInfo {
print_generated_seed(&args.seed, true);
}
let mut result = Ok(());
for round in args.start_round..args.rounds {
if args.rounds > 1 {
let tod = Local::now().format("%F %R");
let end = if args.rounds == u64::MAX {
"inf]".to_string()
} else {
format!("{})", args.rounds)
};
println!(
"{}[{}] Round {} in range [{}, {} ...",
if round > args.start_round { "\n" } else { "" },
tod,
round,
args.start_round,
end,
);
}
let round_id = if args.write {
round
} else {
args.start_round
};
result = Ok(());
if args.write {
let (mut disktest, file) = new_disktest(&args, round_id, true, &abort)?;
result = disktest.write(file, args.seek, args.max_bytes).map(|_| ());
}
if args.verify && result.is_ok() {
let (mut disktest, file) = new_disktest(&args, round_id, false, &abort)?;
result = disktest.verify(file, args.seek, args.max_bytes).map(|_| ());
}
if result.is_err() {
break;
}
}
if !args.user_seed && args.quiet < DisktestQuiet::NoInfo {
print_generated_seed(&args.seed, false);
}
if result.is_ok() && args.quiet == DisktestQuiet::Normal {
println!("Success!");
}
result
}