提交 e1c22265 编写于 作者: G gongzhengyang

feat: add example subcommands

上级 fce36a47
use clap::{arg, ArgMatches, Command};
use clap::{Arg, ArgMatches, Command};
pub fn get_matches() -> ArgMatches {
let db_arg = Arg::new("db")
.long("db")
.help("the xdb filepath, you can set this field like \
../data/ip2region.xdb,if you dont set,\
if will detect xdb file on ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb if exists");
Command::new("ip2region")
.version("0.1")
.about("ip2region bin program")
.long_about(
"you can set environment XDB_FILEPATH=../data/ip2region or just use --xdb in command",
)
.arg(
arg!(--xdb <xdb> "the xdb filepath, you can set this field like ../data/ip2region.xdb"),
.long_about("you can set --db in command to specific the xdb filepath, default run query")
.subcommand(Command::new("query").about("query test").arg(&db_arg))
.subcommand(
Command::new("bench")
.about("bench test")
.arg(
Arg::new("src")
.long("src")
.help("set this to specific source bench file")
.required(true),
)
.arg(&db_arg),
)
.get_matches()
}
extern crate core;
use clap::ArgMatches;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::net::Ipv4Addr;
use std::str::FromStr;
use std::time::Instant;
use ip2region2::{searcher_init, search_by_ip};
use ip2region2::{search_by_ip, searcher_init};
mod cmd;
fn main() {
/// set rust log level
/// set rust log level, if you don`t want print log, you can skip this
fn log_init() {
let rust_log_key = "RUST_LOG";
std::env::var(rust_log_key).unwrap_or_else(|_| {
std::env::set_var(rust_log_key, "INFO");
std::env::var(rust_log_key).unwrap()
});
tracing_subscriber::fmt::init();
}
/// init default xdb_filepath config
/// if value if None, if will detect xdb file on ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb if exists
let matches = cmd::get_matches();
if let Some(xdb_filepath) = matches.get_one::<String>("xdb") {
searcher_init(Some(xdb_filepath.to_owned()))
} else {
searcher_init(None);
fn bench_test(src_filepath: &str) {
let now = Instant::now();
let mut count = 0;
let mut file = File::open(src_filepath).unwrap();
let mut contents = String::new();
file.read_to_string(&mut contents).unwrap();
for line in contents.split('\n') {
if !line.contains('|') {
continue;
}
let ip_test_line = line.splitn(3, '|').collect::<Vec<&str>>();
let start_ip = Ipv4Addr::from_str(ip_test_line[0]).unwrap();
let end_ip = Ipv4Addr::from_str(ip_test_line[1]).unwrap();
if end_ip < start_ip {
panic!("start ip({start_ip}) should not be greater than end ip({end_ip})")
}
let start_ip = u32::from(start_ip);
let end_ip = u32::from(end_ip);
let mid_ip = (((start_ip as u64) + (end_ip as u64)) >> 1) as u32;
for ip in [
start_ip,
((start_ip as u64 + mid_ip as u64) >> 1) as u32,
mid_ip,
((mid_ip as u64 + end_ip as u64) >> 1) as u32,
end_ip,
] {
search_by_ip(ip).unwrap();
count += 1;
}
}
println!(
"Bench finished, total: {count},\
took: {:?} ,\
cost: {:?}/op",
now.elapsed(),
now.elapsed() / count
)
}
fn query_test() {
println!("ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit");
loop {
print!("ip2region>> ");
......@@ -33,8 +71,34 @@ fn main() {
if line.contains("quit") {
break;
}
let line = line.trim();
let now = Instant::now();
let result = search_by_ip(line.trim());
println!("region: {:?}, took: {:?}", result, now.elapsed());
let result = search_by_ip(line);
let cost = now.elapsed();
println!("region: {result:?}, took: {cost:?}", );
}
}
fn matches_for_searcher(matches: &ArgMatches) {
if let Some(xdb_filepath) = matches.get_one::<String>("db") {
searcher_init(Some(xdb_filepath.to_owned()))
} else {
searcher_init(None);
}
}
fn main() {
log_init();
let matches = cmd::get_matches();
if let Some(sub_matches) = matches.subcommand_matches("bench") {
matches_for_searcher(sub_matches);
let src_filepath = sub_matches.get_one::<String>("src").unwrap();
bench_test(src_filepath);
}
if let Some(sub_matches) = matches.subcommand_matches("query") {
matches_for_searcher(sub_matches);
query_test()
}
}
......@@ -2,8 +2,8 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion};
use rand;
use ip2region2::searcher::{
get_block_by_size, get_full_cache, get_vector_index_cache,
search_by_ip, searcher_init, get_int_block_value
get_block_by_size, get_full_cache, get_int_block_value, get_vector_index_cache, search_by_ip,
searcher_init,
};
fn ip_search_bench(c: &mut Criterion) {
......@@ -18,9 +18,11 @@ fn ip_search_bench(c: &mut Criterion) {
fn get_block_by_size_bench(c: &mut Criterion) {
c.bench_function("get_block_by_size", |b| {
b.iter(|| {
black_box(get_block_by_size(get_full_cache(),
rand::random::<u16>() as usize,
4));
black_box(get_block_by_size(
get_full_cache(),
rand::random::<u16>() as usize,
4,
));
})
});
}
......@@ -28,8 +30,10 @@ fn get_block_by_size_bench(c: &mut Criterion) {
fn get_int_block_bench(c: &mut Criterion) {
c.bench_function("get_int_block_bench", |b| {
b.iter(|| {
black_box(get_int_block_value(get_full_cache(),
rand::random::<u16>() as usize));
black_box(get_int_block_value(
get_full_cache(),
rand::random::<u16>() as usize,
));
})
});
}
......
......@@ -81,8 +81,7 @@ pub fn get_block_by_size(bytes: &[u8], offset: usize, length: usize) -> usize {
result
}
pub fn searcher_init(xdb_filepath: Option<String>)
{
pub fn searcher_init(xdb_filepath: Option<String>) {
let xdb_filepath = xdb_filepath.unwrap_or_else(|| default_detect_xdb_file().unwrap());
std::env::set_var(XDB_FILEPATH_ENV, xdb_filepath);
CACHE.get_or_init(load_file);
......@@ -159,4 +158,16 @@ mod tests {
println!("ip search in main thread: {r}");
handle.join().unwrap();
}
#[test]
fn test_multi_searcher_init() {
for _ in 0..5 {
thread::spawn(|| {
searcher_init(None);
});
}
searcher_init(None);
searcher_init(Some(String::from("test")));
search_by_ip(123).unwrap();
}
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册