xinb
/
ip2region_private
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: change searcher fn

This commit is contained in:
gongzhengyang 2022-12-23 17:06:14 +08:00
parent 7e7a612437
commit 654080b1d0
6 changed files with 84 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
binding/rust/example/Cargo.toml
View File

@ -12,3 +12,5 @@ license = "Apache-2.0"
[dependencies]
ip2region2 = { path = "../ip2region2" }
clap = { version = "4.0" }
tracing = "0.1"
tracing-subscriber = "0.2.0"

28
binding/rust/example/src/main.rs
View File

@ -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::<String>("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::<String>("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());
}
}

2
binding/rust/ip2region2/Cargo.toml
View File

@ -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"

54
binding/rust/ip2region2/benches/search.rs
View File

@ -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::<u32>()).unwrap();
})
});
}
fn buffer_value_bench(c: &mut Criterion) {
c.bench_function("buffer_value", |b| {
b.iter(|| {
let offset = rand::random::<u16>();
let length = rand::random::<u8>();
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::<u16>() as usize,
4,
);
black_box(get_block_by_size(get_full_cache(),
rand::random::<u16>() 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::<u32>());
black_box(get_int_block_value(get_full_cache(),
rand::random::<u16>() 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);

3
binding/rust/ip2region2/src/lib.rs
View File

@ -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};

81
binding/rust/ip2region2/src/searcher.rs
View File

@ -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<Vec<u8>> = OnceCell::new();
/// check https://mp.weixin.qq.com/s/ndjzu0BgaeBmDOCw5aqHUg for details
pub fn search_by_ip<T>(ip: T) -> Result<String, Box<dyn Error>>
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<T>(ip: T) -> Result<String, Box<dyn Error>>
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<T>(ip: T) -> Result<String, Box<dyn Error>>
} 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<String, Box<dyn Error>> {
}
#[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<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);
}
pub fn searcher_init(xdb_filepath: Option<String>, cache_policy: Option<CachePolicy>) {
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<u8> = get_full_cache();
&full_cache[HEADER_INFO_LENGTH..(HEADER_INFO_LENGTH + VECTOR_INDEX_LENGTH)]
}
fn load_file() -> Vec<u8>{
let xdb_filepath = std::env::var("XDB_FILEPATH").unwrap();
fn load_file() -> Vec<u8> {
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<u8>{
buffer
}
fn get_full_cache() -> &'static Vec<u8> {
let cache_policy = std::env::var(CACHE_POLICY_ENV).unwrap();
if cache_policy == CachePolicy::Full {
static CACHE: OnceCell<Vec<u8>> = OnceCell::new();
return CACHE.get_or_init(|| load_file())
}
&load_file()
pub fn get_full_cache() -> &'static Vec<u8> {
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();