summaryrefslogtreecommitdiffstats
path: root/disktest/src/main.rs
blob: 44e9de0ee346bb5b90e0e0d3463ffd4855e8f078 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
// -*- coding: utf-8 -*-
//
// disktest - Storage tester
//
// Copyright 2020-2024 Michael Büsch <m@bues.ch>
//
// Licensed under the Apache License version 2.0
// or the MIT license, at your option.
// SPDX-License-Identifier: Apache-2.0 OR MIT
//

#![forbid(unsafe_code)]

mod args;

use anyhow as ah;
use args::{parse_args, Args};
use chrono::prelude::*;
use disktest_lib::{Disktest, DisktestFile, DisktestQuiet};
use std::{
    env::args_os,
    sync::{atomic::AtomicBool, Arc},
};

/// Install abort signal handlers and return
/// the abort-flag that is written to true by these handlers.
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)
}

/// Print the generated seed to the console.
fn print_generated_seed(seed: &str, verbose: bool) {
    if verbose {
        println!(
            "\nThe generated --seed is:\n    {}\nUse this seed for subsequent --verify.\n",
            seed
        );
    } else {
        println!("Generated --seed {}\n", seed);
    }
}

/// Create a new disktest core instance.
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)?,
    ))
}

/// Main program entry point.
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 {
            // In verify-only mode it doesn't make sense to change the key per round.
            args.start_round
        };

        // Run write-mode, if requested.
        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(|_| ());
        }

        // Run verify-mode, if requested.
        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
}

// vim: ts=4 sw=4 expandtab
bues.ch cgit interface