kvbench/
cmdline.rs

1use crate::stores::Registry;
2use clap::ValueHint::FilePath;
3use clap::{Args, Parser, Subcommand};
4use log::debug;
5use std::fs::read_to_string;
6use std::sync::mpsc::channel;
7
8#[derive(Args, Debug)]
9struct BenchArgs {
10    #[arg(short = 's')]
11    #[arg(value_hint = FilePath)]
12    #[arg(help = "Path to the key-value store's TOML config file")]
13    store_config: String,
14
15    #[arg(short = 'b')]
16    #[arg(value_hint = FilePath)]
17    #[arg(help = "Path to the benchmark's TOML config file")]
18    benchmark_config: String,
19}
20
21#[derive(Args, Debug)]
22struct ServerArgs {
23    #[arg(short = 'a', default_value = "0.0.0.0")]
24    #[arg(help = "Bind address")]
25    host: String,
26
27    #[arg(short = 'p', default_value = "9000")]
28    #[arg(help = "Bind port")]
29    port: String,
30
31    #[arg(short = 's')]
32    #[arg(value_hint = FilePath)]
33    #[arg(help = "Path to the key-value store's TOML config file")]
34    store_config: String,
35
36    #[arg(short = 'n', default_value_t = 1)]
37    #[arg(help = "Number of worker threads")]
38    workers: usize,
39}
40
41#[derive(Parser, Debug)]
42#[command(version, about)]
43struct Cli {
44    #[command(subcommand)]
45    command: Commands,
46}
47
48#[derive(Subcommand, Debug)]
49enum Commands {
50    #[command(about = "Run a benchmark")]
51    Bench(BenchArgs),
52    #[command(about = "Start a key-value server")]
53    Server(ServerArgs),
54    #[command(about = "List all registered key-value stores")]
55    List,
56}
57
58fn bench_cli(args: &BenchArgs) {
59    let opt: String = {
60        let s = args.store_config.clone();
61        let b = args.benchmark_config.clone();
62        read_to_string(s.as_str()).unwrap() + "\n" + &read_to_string(b.as_str()).unwrap()
63    };
64
65    let (map, phases) = crate::bench::init(&opt);
66    map.bench(&phases);
67}
68
69fn server_cli(args: &ServerArgs) {
70    let host = &args.host;
71    let port = &args.port;
72    let nr_workers = args.workers;
73
74    let opt: String = read_to_string(args.store_config.as_str()).unwrap();
75    let map = crate::server::init(&opt);
76
77    let (stop_tx, stop_rx) = channel();
78    let (grace_tx, grace_rx) = channel();
79
80    ctrlc::set_handler(move || {
81        assert!(stop_tx.send(()).is_ok());
82        debug!("SIGINT received and stop message sent to server");
83    })
84    .expect("Error setting Ctrl-C handler for server");
85
86    map.server(&host, &port, nr_workers, stop_rx, grace_tx);
87
88    assert!(grace_rx.recv().is_ok());
89    debug!("All server threads have been shut down gracefully, exit");
90}
91
92fn list_cli() {
93    for r in inventory::iter::<Registry> {
94        println!("Registered map: {}", r.name);
95    }
96}
97
98/// The default command line interface.
99///
100/// This function is public and can be called in a different crate. For example, one can integrate
101/// their own key-value stores by registering the constructor function. Then, adding this function
102/// will produce a benchmark binary the has the same usage as the one in this crate.
103///
104/// ## Usage
105///
106/// To get the usage of the command line interface, users can run:
107///
108/// ```bash
109/// kvbench -h
110/// ```
111///
112/// The interface supports three modes, `bench`, `server` and `list`.
113///
114/// ### Benchmark Mode
115///
116/// Usage:
117///
118/// ```bash
119/// kvbench bench -s <STORE_CONFIG> -b <BENCH_CONFIG>
120/// ```
121///
122/// Where `STORE_CONFIG` and `BENCH_CONFIG` are the paths to the key-value store and benchmark
123/// configuration files, respectively.
124///
125/// ### Server mode
126///
127/// Usage:
128///
129/// ```bash
130/// kvbench server -s <STORE_CONFIG> -a <HOST> -p <PORT> -n <WORKERS>
131/// ```
132///
133/// Where `STORE_CONFIG` is the path of the key-value store configuration file.
134///
135/// The default `HOST` and `PORT` are `0.0.0.0` and `9000`. By default, the server will spawn one
136/// worker thread only for incoming connections. You can adjust the number of worker threads by
137/// specifying `-n`.
138///
139/// ### List mode
140///
141/// Usage:
142/// ``` bash
143/// kvbench list
144/// ```
145///
146/// This command lists all registered key-value stores' names.
147
148pub fn cmdline() {
149    env_logger::init();
150    let cli = Cli::parse();
151    debug!("Starting kvbench with args: {:?}", cli);
152    match cli.command {
153        Commands::Bench(args) => bench_cli(&args),
154        Commands::Server(args) => server_cli(&args),
155        Commands::List => list_cli(),
156    }
157}
kvbench/cmdline.rs

kvbench/
cmdline.rs
