diff --git a/binding/rust/example/src/cmd.rs b/binding/rust/example/src/cmd.rs index 05c3912..98402a4 100644 --- a/binding/rust/example/src/cmd.rs +++ b/binding/rust/example/src/cmd.rs @@ -1,14 +1,27 @@ -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 "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() } diff --git a/binding/rust/example/src/main.rs b/binding/rust/example/src/main.rs index 2bfab9d..bf5c874 100644 --- a/binding/rust/example/src/main.rs +++ b/binding/rust/example/src/main.rs @@ -1,29 +1,67 @@ +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::("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::>(); + 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::("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::("src").unwrap(); + + bench_test(src_filepath); + } + + if let Some(sub_matches) = matches.subcommand_matches("query") { + matches_for_searcher(sub_matches); + query_test() } } diff --git a/binding/rust/ip2region2/benches/search.rs b/binding/rust/ip2region2/benches/search.rs index ff8bed1..3854502 100644 --- a/binding/rust/ip2region2/benches/search.rs +++ b/binding/rust/ip2region2/benches/search.rs @@ -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::() as usize, - 4)); + black_box(get_block_by_size( + get_full_cache(), + rand::random::() 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::() as usize)); + black_box(get_int_block_value( + get_full_cache(), + rand::random::() as usize, + )); }) }); } diff --git a/binding/rust/ip2region2/src/searcher.rs b/binding/rust/ip2region2/src/searcher.rs index d31fd6c..f7dc20b 100644 --- a/binding/rust/ip2region2/src/searcher.rs +++ b/binding/rust/ip2region2/src/searcher.rs @@ -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) -{ +pub fn searcher_init(xdb_filepath: Option) { 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(); + } }