diff --git a/binding/rust/ip2region2/Cargo.toml b/binding/rust/ip2region2/Cargo.toml index 6aa3475..11e5a02 100644 --- a/binding/rust/ip2region2/Cargo.toml +++ b/binding/rust/ip2region2/Cargo.toml @@ -10,6 +10,8 @@ license = "Apache-2.0" [dependencies] once_cell = "1.16" +tracing = "0.1" +#tracing-subscriber = "0.2" [dev-dependencies] criterion = "0.4" diff --git a/binding/rust/ip2region2/src/lib.rs b/binding/rust/ip2region2/src/lib.rs index 97f5a95..c307f85 100644 --- a/binding/rust/ip2region2/src/lib.rs +++ b/binding/rust/ip2region2/src/lib.rs @@ -1,121 +1,3 @@ -use std::error::Error; -use std::fmt::Display; -use ip_value::ToUIntIP; - mod ip_value; +pub use self::ip_value::ToUIntIP; mod searcher; - -pub use searcher::global_searcher; - -const HEADER_INFO_LENGTH: usize = 256; -const VECTOR_INDEX_COLS: usize = 256; -const VECTOR_INDEX_SIZE: usize = 8; -const SEGMENT_INDEX_SIZE: usize = 14; - - -pub fn get_start_end_ptr(ip: u32) -> (usize, usize) { - let il0 = ((ip >> 24) & 0xFF) as usize; - let il1 = ((ip >> 16) & 0xFF) as usize; - let idx = VECTOR_INDEX_SIZE * (il0 * VECTOR_INDEX_COLS + il1); - let start_point = HEADER_INFO_LENGTH + idx; - - let start_ptr = get_block_by_size(global_searcher().buffer(), start_point, 4); - let end_ptr = get_block_by_size(global_searcher().buffer(), start_point + 4, 4); - (start_ptr, end_ptr) -} - -/// check https://mp.weixin.qq.com/s/ndjzu0BgaeBmDOCw5aqHUg for details -pub fn search_by_ip(ip: T) -> Result> -where - T: ToUIntIP + Display, -{ - let ip = ip.to_u32_ip()?; - let (start_ptr, end_ptr) = get_start_end_ptr(ip); - let mut left: usize = 0; - let mut right: usize = (end_ptr - start_ptr) / SEGMENT_INDEX_SIZE; - - while left <= right { - let mid = (left + right) >> 1; - let offset = start_ptr + mid * SEGMENT_INDEX_SIZE; - let buffer_ip_value = buffer_value(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; - } else if ip > (get_block_by_size(buffer_ip_value, 4, 4) as u32) { - left = mid + 1; - } 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(buffer_value(data_offset, data_length).to_vec()); - return Ok(result?); - } - } - Err("not matched".into()) -} - -pub fn buffer_value(offset: usize, length: usize) -> &'static [u8] { - &global_searcher().buffer()[offset..offset + length] -} - -#[inline] -pub fn get_block_by_size(bytes: &[T], offset: usize, length: usize) -> usize -where - T: Clone, - usize: From, -{ - let mut result: usize = 0; - for (index, value) in bytes[offset..offset + length].iter().enumerate() { - result |= usize::from(value.clone()) << (index * 8); - } - result -} - -#[cfg(test)] -mod tests { - 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() { - search_by_ip("2.0.0.0").unwrap(); - search_by_ip("32").unwrap(); - search_by_ip(4294408949).unwrap(); - search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap(); - } - - #[test] - fn test_match_all_ip_correct() { - let mut file = File::open("../../../data/ip.test.txt").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(); - for value in u32::from(start_ip)..u32::from(end_ip) + 1 { - let result = search_by_ip(value).unwrap(); - assert_eq!(result.as_str(), ip_test_line[2]) - } - } - } - - #[test] - fn test_multi_thread_only_load_xdb_once() { - let handle = thread::spawn(|| { - 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(); - println!("ip search in main thread: {r}"); - handle.join().unwrap(); - } -} diff --git a/binding/rust/ip2region2/src/searcher.rs b/binding/rust/ip2region2/src/searcher.rs index abb64d9..8e4b4a6 100644 --- a/binding/rust/ip2region2/src/searcher.rs +++ b/binding/rust/ip2region2/src/searcher.rs @@ -1,4 +1,3 @@ -use std::env; use std::error::Error; use std::fmt; use std::fmt::{Display, Formatter}; @@ -8,59 +7,181 @@ use std::path::Path; use once_cell::sync::OnceCell; +use crate::ToUIntIP; + +const HEADER_INFO_LENGTH: usize = 256; +const VECTOR_INDEX_COLS: usize = 256; +const VECTOR_INDEX_SIZE: usize = 8; +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, + Never=1, VecIndex, Full, } -/// store the xdb file in memory totally -pub struct Searcher { - vec_cache: Vec, - full_cache: Vec, +/// check https://mp.weixin.qq.com/s/ndjzu0BgaeBmDOCw5aqHUg for details +pub fn search_by_ip(ip: T) -> Result> + where + T: ToUIntIP + Display, +{ + let ip = ip.to_u32_ip()?; + let (start_ptr, end_ptr) = get_start_end_ptr(ip); + let mut left: usize = 0; + let mut right: usize = (end_ptr - start_ptr) / SEGMENT_INDEX_SIZE; + + 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 start_ip = get_block_by_size(buffer_ip_value, 0, 4); + if ip < (start_ip as u32) { + right = mid - 1; + } else if ip > (get_block_by_size(buffer_ip_value, 4, 4) as u32) { + left = mid + 1; + } 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()); + return Ok(result?); + } + } + Err("not matched".into()) } -impl Searcher { - pub fn new(xdb_filepath: Option<&str>, cache_policy: Option) -> Result> { - let xdb_filepath = xdb_filepath.unwrap_or_else(|_| { - Searcher::default_detect_xdb_file().unwrap().as_str() - }); - println!("load xdb searcher file at {xdb_filepath}"); - let mut f = File::open(xdb_filepath)?; - let mut buffer = Vec::new(); - f.read_to_end(&mut buffer)?; - Ok(Self { buffer }) +pub fn get_start_end_ptr(ip: u32) -> (usize, usize) { + let il0 = ((ip >> 24) & 0xFF) as usize; + let il1 = ((ip >> 16) & 0xFF) as 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 end_ptr = get_block_by_size(vector_cache, start_point + 4, 4); + (start_ptr, end_ptr) +} + +/// it will check ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb +fn default_detect_xdb_file() -> Result> { + let prefix = "../".to_owned(); + for recurse in 1..4 { + let filepath = prefix.repeat(recurse) + "data/ip2region.xdb"; + if Path::new(filepath.as_str()).exists() { + return Ok(filepath); + } + } + Err("default filepath not find the xdb file, so you must set xdb_filepath".into()) +} + +#[inline] +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 +} + +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, 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] { + 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(); + 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(); + f.read_to_end(&mut buffer).expect("load file error"); + 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() +} + +#[cfg(test)] +mod tests { + 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); + + search_by_ip("2.0.0.0").unwrap(); + search_by_ip("32").unwrap(); + search_by_ip(4294408949).unwrap(); + search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap(); } - /// it will check ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb - fn default_detect_xdb_file() -> Result> { - let prefix = "../".to_owned(); - for recurse in 1..4 { - let filepath = prefix.repeat(recurse) + "data/ip2region.xdb"; - if Path::new(filepath.as_str()).exists() { - return Ok(filepath); + #[test] + fn test_match_all_ip_correct() { + searcher_init(None, None); + let mut file = File::open("../../../data/ip.test.txt").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(); + for value in u32::from(start_ip)..u32::from(end_ip) + 1 { + let result = search_by_ip(value).unwrap(); + assert_eq!(result.as_str(), ip_test_line[2]) } } - Err("default filepath not find the xdb file, so you must set xdb_filepath".into()) } - pub fn buffer(&self) -> &Vec { - self.full_cache.as_ref() - } - - pub fn vec_cache(&self) -> &Vec { - self.vec_cache.as_ref() - } -} - -/// global init searcher thread safely -pub fn global_searcher() -> &'static Searcher { - static SEARCHER: OnceCell = OnceCell::new(); - SEARCHER.get_or_init(|| Searcher::new().unwrap()) -} - -impl Display for Searcher { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "searcher_with_len {}", self.buffer.len()) + #[test] + fn test_multi_thread_only_load_xdb_once() { + searcher_init(None, None); + let handle = thread::spawn(|| { + 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(); + println!("ip search in main thread: {r}"); + handle.join().unwrap(); } }