From 654080b1d06470a42d21c2d8632bd0af17b9d456 Mon Sep 17 00:00:00 2001 From: gongzhengyang Date: Fri, 23 Dec 2022 17:06:14 +0800 Subject: [PATCH] feat: change searcher fn --- binding/rust/example/Cargo.toml | 2 + binding/rust/example/src/main.rs | 28 +++++--- binding/rust/ip2region2/Cargo.toml | 2 +- binding/rust/ip2region2/benches/search.rs | 54 +++++++++------ binding/rust/ip2region2/src/lib.rs | 3 +- binding/rust/ip2region2/src/searcher.rs | 81 ++++++++--------------- 6 files changed, 84 insertions(+), 86 deletions(-) diff --git a/binding/rust/example/Cargo.toml b/binding/rust/example/Cargo.toml index b19caf7..3d0072f 100644 --- a/binding/rust/example/Cargo.toml +++ b/binding/rust/example/Cargo.toml @@ -12,3 +12,5 @@ license = "Apache-2.0" [dependencies] ip2region2 = { path = "../ip2region2" } clap = { version = "4.0" } +tracing = "0.1" +tracing-subscriber = "0.2.0" diff --git a/binding/rust/example/src/main.rs b/binding/rust/example/src/main.rs index 4d2b97f..2bfab9d 100644 --- a/binding/rust/example/src/main.rs +++ b/binding/rust/example/src/main.rs @@ -1,19 +1,29 @@ -use std::env; use std::io::Write; use std::time::Instant; +use ip2region2::{searcher_init, search_by_ip}; + mod cmd; fn main() { - env::var("XDB_FILEPATH").unwrap_or_else(|_| { - let matches = cmd::get_matches(); - if let Some(xdb_filepath) = matches.get_one::("xdb") { - env::set_var("XDB_FILEPATH", xdb_filepath); - } - "".to_owned() + /// set rust log level + 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); + } + - ip2region2::global_searcher(); println!("ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit"); loop { print!("ip2region>> "); @@ -24,7 +34,7 @@ fn main() { break; } let now = Instant::now(); - let result = ip2region2::search_by_ip(line.trim()); + let result = search_by_ip(line.trim()); println!("region: {:?}, took: {:?}", result, now.elapsed()); } } diff --git a/binding/rust/ip2region2/Cargo.toml b/binding/rust/ip2region2/Cargo.toml index 11e5a02..a4d51bb 100644 --- a/binding/rust/ip2region2/Cargo.toml +++ b/binding/rust/ip2region2/Cargo.toml @@ -11,7 +11,7 @@ license = "Apache-2.0" [dependencies] once_cell = "1.16" tracing = "0.1" -#tracing-subscriber = "0.2" +tracing-subscriber = "0.2.0" [dev-dependencies] criterion = "0.4" diff --git a/binding/rust/ip2region2/benches/search.rs b/binding/rust/ip2region2/benches/search.rs index b3a8411..ff8bed1 100644 --- a/binding/rust/ip2region2/benches/search.rs +++ b/binding/rust/ip2region2/benches/search.rs @@ -1,42 +1,51 @@ -use criterion::{criterion_group, criterion_main, Criterion}; +use criterion::{black_box, criterion_group, criterion_main, Criterion}; use rand; -use ip2region2::{buffer_value, get_block_by_size, get_start_end_ptr, global_searcher, search_by_ip}; +use ip2region2::searcher::{ + get_block_by_size, get_full_cache, get_vector_index_cache, + search_by_ip, searcher_init, get_int_block_value +}; fn ip_search_bench(c: &mut Criterion) { c.bench_function("ip_search_bench", |b| { + searcher_init(None); b.iter(|| { search_by_ip(rand::random::()).unwrap(); }) }); } -fn buffer_value_bench(c: &mut Criterion) { - c.bench_function("buffer_value", |b| { - b.iter(|| { - let offset = rand::random::(); - let length = rand::random::(); - buffer_value(offset as usize, length as usize); - }); - }); -} - fn get_block_by_size_bench(c: &mut Criterion) { c.bench_function("get_block_by_size", |b| { b.iter(|| { - get_block_by_size( - &global_searcher().buffer(), - rand::random::() as usize, - 4, - ); + black_box(get_block_by_size(get_full_cache(), + rand::random::() as usize, + 4)); }) }); } -fn get_start_end_ptr_bench(c: &mut Criterion) { - c.bench_function("get_start_end_ptr", |b| { +fn get_int_block_bench(c: &mut Criterion) { + c.bench_function("get_int_block_bench", |b| { b.iter(|| { - get_start_end_ptr(rand::random::()); + black_box(get_int_block_value(get_full_cache(), + rand::random::() as usize)); + }) + }); +} + +fn get_full_cache_bench(c: &mut Criterion) { + c.bench_function("get_full_cache", |b| { + b.iter(|| { + black_box(get_full_cache()); + }) + }); +} + +fn get_vec_index_cache_bench(c: &mut Criterion) { + c.bench_function("get_vec_index_cache", |b| { + b.iter(|| { + black_box(get_vector_index_cache()); }) }); } @@ -44,8 +53,9 @@ fn get_start_end_ptr_bench(c: &mut Criterion) { criterion_group!( benches, ip_search_bench, - buffer_value_bench, + get_int_block_bench, get_block_by_size_bench, - get_start_end_ptr_bench + get_full_cache_bench, + get_vec_index_cache_bench, ); criterion_main!(benches); diff --git a/binding/rust/ip2region2/src/lib.rs b/binding/rust/ip2region2/src/lib.rs index c307f85..350f1f2 100644 --- a/binding/rust/ip2region2/src/lib.rs +++ b/binding/rust/ip2region2/src/lib.rs @@ -1,3 +1,4 @@ mod ip_value; pub use self::ip_value::ToUIntIP; -mod searcher; +pub mod searcher; +pub use searcher::{search_by_ip, searcher_init}; diff --git a/binding/rust/ip2region2/src/searcher.rs b/binding/rust/ip2region2/src/searcher.rs index 8e4b4a6..d31fd6c 100644 --- a/binding/rust/ip2region2/src/searcher.rs +++ b/binding/rust/ip2region2/src/searcher.rs @@ -1,6 +1,5 @@ use std::error::Error; -use std::fmt; -use std::fmt::{Display, Formatter}; +use std::fmt::Display; use std::fs::File; use std::io::Read; use std::path::Path; @@ -16,19 +15,13 @@ const SEGMENT_INDEX_SIZE: usize = 14; const VECTOR_INDEX_LENGTH: usize = 512 * 1024; const XDB_FILEPATH_ENV: &str = "XDB_FILEPATH"; -const CACHE_POLICY_ENV: &str = "CACHE_POLICY"; -#[derive(Debug, Copy, Clone, PartialEq)] -pub enum CachePolicy { - Never=1, - VecIndex, - Full, -} +static CACHE: OnceCell> = OnceCell::new(); /// check https://mp.weixin.qq.com/s/ndjzu0BgaeBmDOCw5aqHUg for details pub fn search_by_ip(ip: T) -> Result> - where - T: ToUIntIP + Display, +where + T: ToUIntIP + Display, { let ip = ip.to_u32_ip()?; let (start_ptr, end_ptr) = get_start_end_ptr(ip); @@ -38,7 +31,7 @@ pub fn search_by_ip(ip: T) -> Result> while left <= right { let mid = (left + right) >> 1; let offset = start_ptr + mid * SEGMENT_INDEX_SIZE; - let buffer_ip_value = &get_full_cache()[offset..offset+SEGMENT_INDEX_SIZE]; + let buffer_ip_value = &get_full_cache()[offset..offset + SEGMENT_INDEX_SIZE]; let start_ip = get_block_by_size(buffer_ip_value, 0, 4); if ip < (start_ip as u32) { right = mid - 1; @@ -47,7 +40,9 @@ pub fn search_by_ip(ip: T) -> Result> } else { let data_length = get_block_by_size(buffer_ip_value, 8, 2); let data_offset = get_block_by_size(buffer_ip_value, 10, 4); - let result = String::from_utf8(get_full_cache()[data_offset..(data_offset + data_length)].to_vec()); + let result = String::from_utf8( + get_full_cache()[data_offset..(data_offset + data_length)].to_vec(), + ); return Ok(result?); } } @@ -60,7 +55,7 @@ pub fn get_start_end_ptr(ip: u32) -> (usize, usize) { let idx = VECTOR_INDEX_SIZE * (il0 * VECTOR_INDEX_COLS + il1); let start_point = idx; let vector_cache = get_vector_index_cache(); - let start_ptr = get_block_by_size( vector_cache, start_point, 4); + let start_ptr = get_block_by_size(vector_cache, start_point, 4); let end_ptr = get_block_by_size(vector_cache, start_point + 4, 4); (start_ptr, end_ptr) } @@ -78,44 +73,29 @@ fn default_detect_xdb_file() -> Result> { } #[inline] -pub fn get_block_by_size(bytes: &[u8], offset: usize, length: usize) -> usize -{ +pub fn get_block_by_size(bytes: &[u8], offset: usize, length: usize) -> usize { let mut result: usize = 0; for (index, value) in bytes[offset..offset + length].iter().enumerate() { - result |= usize::from(value.clone()) << (index * 8); + result += usize::from(*value) << (index << 3); } result } -fn set_log_level() { - 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() - }); +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); } -pub fn searcher_init(xdb_filepath: Option, cache_policy: Option) { - set_log_level(); - let xdb_filepath = xdb_filepath.unwrap_or_else(|| { - default_detect_xdb_file().unwrap() - }); - std::env::set_var(XDB_FILEPATH_ENV, xdb_filepath.as_str()); - if let Some(policy) = cache_policy { - std::env::set_var(CACHE_POLICY_ENV, policy); - return; - } - std::env::set_var(CACHE_POLICY_ENV, CachePolicy::Full); - -} - -fn get_vector_index_cache() -> &'static [u8] { +pub fn get_vector_index_cache() -> &'static [u8] { let full_cache: &'static Vec = get_full_cache(); &full_cache[HEADER_INFO_LENGTH..(HEADER_INFO_LENGTH + VECTOR_INDEX_LENGTH)] } -fn load_file() -> Vec{ - let xdb_filepath = std::env::var("XDB_FILEPATH").unwrap(); +fn load_file() -> Vec { + let xdb_filepath = + std::env::var("XDB_FILEPATH").unwrap_or_else(|_| default_detect_xdb_file().unwrap()); tracing::debug!("load xdb searcher file at {} ", xdb_filepath); let mut f = File::open(xdb_filepath).expect("file open error"); let mut buffer = Vec::new(); @@ -123,29 +103,24 @@ fn load_file() -> Vec{ buffer } -fn get_full_cache() -> &'static Vec { - let cache_policy = std::env::var(CACHE_POLICY_ENV).unwrap(); - if cache_policy == CachePolicy::Full { - static CACHE: OnceCell> = OnceCell::new(); - return CACHE.get_or_init(|| load_file()) - } - &load_file() +pub fn get_full_cache() -> &'static Vec { + CACHE.get_or_init(load_file) } #[cfg(test)] mod tests { + use std::fs::File; + use std::io::Read; use std::net::Ipv4Addr; use std::str::FromStr; use std::thread; - use std::fs::File; - use std::io::Read; use super::*; ///test all types find correct #[test] fn test_multi_type_ip() { - searcher_init(None, None); + searcher_init(None); search_by_ip("2.0.0.0").unwrap(); search_by_ip("32").unwrap(); @@ -155,7 +130,7 @@ mod tests { #[test] fn test_match_all_ip_correct() { - searcher_init(None, None); + searcher_init(None); let mut file = File::open("../../../data/ip.test.txt").unwrap(); let mut contents = String::new(); file.read_to_string(&mut contents).unwrap(); @@ -175,9 +150,9 @@ mod tests { #[test] fn test_multi_thread_only_load_xdb_once() { - searcher_init(None, None); + searcher_init(None); let handle = thread::spawn(|| { - let result =search_by_ip("2.2.2.2").unwrap(); + let result = search_by_ip("2.2.2.2").unwrap(); println!("ip search in spawn: {result}"); }); let r = search_by_ip("1.1.1.1").unwrap();