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

//! A command for running the speed test suite.

use std::fmt::Debug;
use std::hash::BuildHasher;
use std::num::NonZeroUsize;

use anyhow::anyhow;
use bitbelay_providers::Provider;
use bitbelay_report::Config;
use bitbelay_suites::performance;
use bitbelay_suites::r#trait::Suite as _;
use byte_unit::Byte;

/// Arguments for the avalanche command.
#[derive(clap::Args, Debug)]
pub struct Args {
    /// The amount of data to generate.
    #[clap(short, long, default_value = "10 MB")]
    data_size: String,

    /// The number of iterations to complete.
    #[arg(short, long, default_value_t = 1 << 8)]
    iterations: usize,

    /// The threshold needed for the speed test to pass in megabytes per second.
    #[arg(short, long, default_value_t = 1000.0)]
    threshold: f64,
}

/// The main function for the speed command.
pub fn main<H: BuildHasher>(
    args: Args,
    build_hasher: H,
    provider: Box<dyn Provider>,
) -> anyhow::Result<()> {
    tracing::info!("Starting speed test suite.");

    let iterations = NonZeroUsize::try_from(args.iterations)
        .map_err(|_| anyhow!("--iterations must be non-zero!"))?;

    let desired_data_size = args
        .data_size
        .parse::<Byte>()
        .map_err(|_| anyhow!("invalid data size: {}", args.data_size))?;

    let mut suite = performance::suite::Builder::default()
        .build_hasher(&build_hasher)
        .unwrap()
        .try_build()
        .unwrap();

    suite.run_speed_test(provider, iterations, desired_data_size, args.threshold)?;

    suite
        .report()
        .write_to(&mut std::io::stderr(), &Config::default())?;

    Ok(())
}