bench/bin: Support simple command line args
#![allow(needless_pass_by_value)]
#![allow(unreadable_literal)]
+extern crate clap;
extern crate futures;
extern crate grpcio as grpc;
extern crate kvproto;
@@ -57,6 +58,7 @@ mod pd;
mod transport_simulate;
use test::BenchSamples;
+use clap::{App, Arg, ArgGroup};
/// shortcut to bench a function.
macro_rules! bench {
@@ -94,7 +96,34 @@ fn main() {
e
);
}
- // TODO allow user to specify flag to just bench some cases.
- raftstore::bench_raftstore();
- mvcc::bench_engine();
+
+ let available_benches = ["raftstore", "mvcc"];
+
+ let matches = App::new("TiKV Benchmark")
+ .args(&available_benches
+ .iter()
+ .map(|name| Arg::with_name(name))
+ .collect::<Vec<_>>())
+ .group(
+ ArgGroup::with_name("benches")
+ .args(&available_benches)
+ .multiple(true),
+ )
+ .get_matches();
+
+ let benches: Vec<_> = if let Some(args) = matches.values_of("benches") {
+ args.collect()
+ } else {
+ available_benches.to_vec()
+ };
+
+ println!("Begin to run: {}", benches.join(", "));
+
+ for item in benches {
+ match item {
+ "raftstore" => raftstore::bench_raftstore(),
+ "mvcc" => mvcc::bench_engine(),
+ _ => eprintln!("error: Unknown bench item {}", item),
+ }
+ }
}
