perf(add more benchmark and change function calls):
This commit is contained in:
parent
e57adde491
commit
76345e67b0
|
|
@ -8,10 +8,11 @@ fn main() {
|
||||||
env::var("XDB_FILEPATH").unwrap_or_else(|_| {
|
env::var("XDB_FILEPATH").unwrap_or_else(|_| {
|
||||||
let matches = cmd::get_matches();
|
let matches = cmd::get_matches();
|
||||||
let xdb_filepath = matches
|
let xdb_filepath = matches
|
||||||
.get_one::<String>("xdb")
|
.get_one::<String>("xdb");
|
||||||
.expect("you must use --xdb in command or set XDB_FILEPATH environment");
|
if xdb_filepath.is_some() {
|
||||||
env::set_var("XDB_FILEPATH", xdb_filepath);
|
env::set_var("XDB_FILEPATH", xdb_filepath.unwrap());
|
||||||
xdb_filepath.to_owned()
|
}
|
||||||
|
"".to_owned()
|
||||||
});
|
});
|
||||||
|
|
||||||
search::global_searcher();
|
search::global_searcher();
|
||||||
|
|
|
||||||
|
|
@ -1,16 +1,41 @@
|
||||||
use criterion::{criterion_group, criterion_main, Criterion};
|
use criterion::{black_box, criterion_group, criterion_main, Criterion};
|
||||||
use rand;
|
use rand;
|
||||||
|
|
||||||
use search::search_by_ip;
|
use search::{buffer_value, get_block_by_size, get_start_end_ptr, global_searcher, search_by_ip};
|
||||||
|
|
||||||
fn ip_search_benchmark(c: &mut Criterion) {
|
fn ip_search_bench(c: &mut Criterion) {
|
||||||
c.bench_function("ip_search_bench", |b| {
|
c.bench_function("ip_search_bench", |b| {
|
||||||
b.iter(|| {
|
b.iter(|| {
|
||||||
let ip = rand::random::<u32>();
|
search_by_ip(rand::random::<u32>()).unwrap();
|
||||||
search_by_ip(ip).unwrap();
|
|
||||||
})
|
})
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
criterion_group!(benches, ip_search_benchmark);
|
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);
|
||||||
|
})
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_start_end_ptr_bench(c: &mut Criterion) {
|
||||||
|
c.bench_function("get_start_end_ptr", |b| {
|
||||||
|
b.iter(|| {
|
||||||
|
get_start_end_ptr(rand::random::<u32>());
|
||||||
|
})
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
criterion_group!(benches, ip_search_bench, buffer_value_bench, get_block_by_size_bench, get_start_end_ptr_bench);
|
||||||
criterion_main!(benches);
|
criterion_main!(benches);
|
||||||
|
|
|
||||||
|
|
@ -1,32 +1,56 @@
|
||||||
mod ip_value;
|
|
||||||
|
|
||||||
use std::env;
|
use std::env;
|
||||||
use std::error::Error;
|
use std::error::Error;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::fmt::Formatter;
|
use std::fmt::Formatter;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
use once_cell::sync::OnceCell;
|
use once_cell::sync::OnceCell;
|
||||||
|
|
||||||
use ip_value::ToUIntIP;
|
use ip_value::ToUIntIP;
|
||||||
|
|
||||||
const HEADER_INFO_LENGTH: u32 = 256;
|
mod ip_value;
|
||||||
// const VECTOR_INDEX_ROWS: u32 = 256;
|
|
||||||
const VECTOR_INDEX_COLS: u32 = 256;
|
const HEADER_INFO_LENGTH: usize = 256;
|
||||||
const VECTOR_INDEX_SIZE: u32 = 8;
|
const VECTOR_INDEX_COLS: usize = 256;
|
||||||
|
const VECTOR_INDEX_SIZE: usize = 8;
|
||||||
const SEGMENT_INDEX_SIZE: usize = 14;
|
const SEGMENT_INDEX_SIZE: usize = 14;
|
||||||
|
|
||||||
|
/// store the xdb file in memory totally
|
||||||
pub struct Searcher {
|
pub struct Searcher {
|
||||||
pub buffer: Vec<u8>,
|
pub buffer: Vec<u8>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl Searcher {
|
||||||
|
/// you can set the XDB_FILEPATH
|
||||||
|
/// or super dir has data dir with the file ip2region.xdb
|
||||||
|
/// it will check ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb
|
||||||
|
pub fn new() -> Result<Self, Box<dyn Error>> {
|
||||||
|
let xdb_filepath = env::var("XDB_FILEPATH")
|
||||||
|
.unwrap_or_else(|_| {
|
||||||
|
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 filepath
|
||||||
|
}
|
||||||
|
};
|
||||||
|
panic!("you must set XDB_FILEPATH or put file in ../data/ip2region.xdb")
|
||||||
|
});
|
||||||
|
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 })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// global init searcher thread safely
|
||||||
pub fn global_searcher() -> &'static Searcher {
|
pub fn global_searcher() -> &'static Searcher {
|
||||||
static SEARCHER: OnceCell<Searcher> = OnceCell::new();
|
static SEARCHER: OnceCell<Searcher> = OnceCell::new();
|
||||||
SEARCHER.get_or_init(|| {
|
SEARCHER.get_or_init(|| {
|
||||||
let xdp_filepath = env::var("XDB_FILEPATH").expect("you must set XDB_FILEPATH for search");
|
Searcher::new().unwrap()
|
||||||
println!("init xdb searcher at {xdp_filepath}");
|
|
||||||
Searcher::new(xdp_filepath.as_str()).unwrap()
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -36,91 +60,89 @@ impl fmt::Display for Searcher {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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<T>(ip: T) -> Result<String, Box<dyn Error>>
|
pub fn search_by_ip<T>(ip: T) -> Result<String, Box<dyn Error>>
|
||||||
where
|
where
|
||||||
T: ToUIntIP,
|
T: ToUIntIP,
|
||||||
{
|
{
|
||||||
let changed_value = ip.to_u32_ip()?;
|
let ip = ip.to_u32_ip()?;
|
||||||
search_by_ip_u32(changed_value)
|
let (start_ptr, end_ptr) = get_start_end_ptr(ip);
|
||||||
}
|
|
||||||
|
|
||||||
pub fn search_by_ip_u32(ip: u32) -> Result<String, Box<dyn Error>> {
|
|
||||||
let il0 = (ip >> 24) & 0xFF;
|
|
||||||
let il1 = (ip >> 16) & 0xFF;
|
|
||||||
let idx = VECTOR_INDEX_SIZE * (il0 * VECTOR_INDEX_COLS + il1);
|
|
||||||
|
|
||||||
let start_point = (HEADER_INFO_LENGTH + idx) as usize;
|
|
||||||
let buffer = &global_searcher().buffer;
|
|
||||||
let start_ptr = get_u32(buffer, start_point);
|
|
||||||
let end_ptr = get_u32(buffer, start_point + 4);
|
|
||||||
let mut left: usize = 0;
|
let mut left: usize = 0;
|
||||||
let mut right: usize = ((end_ptr - start_ptr) as usize) / SEGMENT_INDEX_SIZE;
|
let mut right: usize = (end_ptr - start_ptr) / SEGMENT_INDEX_SIZE;
|
||||||
|
|
||||||
while left <= right {
|
while left <= right {
|
||||||
let mid = (left + right) >> 1;
|
let mid = (left + right) >> 1;
|
||||||
let offset = (start_ptr as usize) + mid * SEGMENT_INDEX_SIZE;
|
let offset = &start_ptr + mid * SEGMENT_INDEX_SIZE;
|
||||||
let buffer_ip_value = buffer_value(offset, SEGMENT_INDEX_SIZE);
|
let buffer_ip_value = buffer_value(offset, SEGMENT_INDEX_SIZE);
|
||||||
let start_ip = get_u32(buffer_ip_value, 0);
|
let start_ip = get_block_by_size(&buffer_ip_value, 0, 4);
|
||||||
if ip < start_ip {
|
if &ip < &(start_ip as u32) {
|
||||||
right = mid - 1;
|
right = mid - 1;
|
||||||
} else if ip > get_u32(buffer_ip_value, 4) {
|
} else if &ip > &(get_block_by_size(&buffer_ip_value, 4, 4) as u32) {
|
||||||
left = mid + 1;
|
left = mid + 1;
|
||||||
} else {
|
} else {
|
||||||
let length = (buffer_ip_value[8] as usize & 0x000000FF)
|
let data_length = get_block_by_size(&buffer_ip_value, 8, 2);
|
||||||
| (buffer_ip_value[9] as usize & 0x0000FF00);
|
let data_offset = get_block_by_size(&buffer_ip_value, 10, 4);
|
||||||
|
let result = String::from_utf8(
|
||||||
let offset = get_u32(buffer_ip_value, 10);
|
buffer_value(data_offset, data_length)
|
||||||
let result = buffer_value(offset as usize, length)
|
.to_vec());
|
||||||
.iter()
|
return Ok(result?);
|
||||||
.map(|x| x.to_owned())
|
|
||||||
.collect::<Vec<u8>>();
|
|
||||||
return Ok(String::from_utf8(result)?);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err("not matched".into())
|
Err("not matched".into())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn start_end_buffer_value(bytes: &[u8], offset: usize, length: usize) -> &[u8] {
|
||||||
|
&bytes[offset..offset+length]
|
||||||
|
}
|
||||||
|
|
||||||
pub fn buffer_value(offset: usize, length: usize) -> &'static [u8] {
|
pub fn buffer_value(offset: usize, length: usize) -> &'static [u8] {
|
||||||
&global_searcher().buffer[offset..offset + length]
|
&global_searcher().buffer[offset..offset + length]
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Searcher {
|
#[inline]
|
||||||
pub fn new(filepath: &str) -> Result<Self, Box<dyn Error>> {
|
pub fn get_block_by_size<T>(bytes: &[T], offset: usize, length: usize) -> usize
|
||||||
let mut f = File::open(filepath)?;
|
where
|
||||||
let mut buffer = Vec::new();
|
T: Clone,
|
||||||
f.read_to_end(&mut buffer)?;
|
usize: From<T>,
|
||||||
Ok(Self { buffer })
|
{
|
||||||
|
let mut result: usize = 0;
|
||||||
|
for (index, value) in bytes[offset..offset+length].iter().enumerate() {
|
||||||
|
result |= usize::from(value.clone()) << (index*8);
|
||||||
}
|
}
|
||||||
}
|
result
|
||||||
|
|
||||||
fn get_u32(bytes: &[u8], offset: usize) -> u32 {
|
|
||||||
(bytes[offset] as u32) & 0x000000FF
|
|
||||||
| ((bytes[offset + 1] as u32) << 8) & 0x0000FF00
|
|
||||||
| ((bytes[offset + 2] as u32) << 16) & 0x00FF0000
|
|
||||||
| ((bytes[offset + 3] as u32) << 24) & 0xFF000000
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
|
||||||
use std::net::Ipv4Addr;
|
use std::net::Ipv4Addr;
|
||||||
use std::str::FromStr;
|
use std::str::FromStr;
|
||||||
use std::thread;
|
use std::thread;
|
||||||
|
|
||||||
const TEST_IP_FILEPATH: &str = "../../../data/ip.test.txt";
|
use super::*;
|
||||||
|
|
||||||
///test all types find correct
|
///test all types find correct
|
||||||
#[test]
|
#[test]
|
||||||
fn test_search_by_ip() {
|
fn test_multi_type_ip() {
|
||||||
search_by_ip("2.0.0.0").unwrap();
|
search_by_ip("2.0.0.0").unwrap();
|
||||||
search_by_ip("32").unwrap();
|
search_by_ip("32").unwrap();
|
||||||
search_by_ip(32).unwrap();
|
search_by_ip(4294408949).unwrap();
|
||||||
search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap();
|
search_by_ip(Ipv4Addr::from_str("1.1.1.1").unwrap()).unwrap();
|
||||||
}
|
}
|
||||||
|
|
||||||
/// test find ip correct use the file ip.test.txt in ../../data
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_random_choose_ip() {
|
fn test_match_all_ip_correct() {
|
||||||
let mut file = File::open(TEST_IP_FILEPATH).unwrap();
|
let mut file = File::open("../../../data/ip.test.txt").unwrap();
|
||||||
let mut contents = String::new();
|
let mut contents = String::new();
|
||||||
file.read_to_string(&mut contents).unwrap();
|
file.read_to_string(&mut contents).unwrap();
|
||||||
for line in contents.split("\n") {
|
for line in contents.split("\n") {
|
||||||
|
|
@ -138,7 +160,7 @@ mod tests {
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_multi_thread() {
|
fn test_multi_thread_only_load_xdb_once() {
|
||||||
let handle = thread::spawn(|| {
|
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}");
|
println!("ip search in spawn: {result}");
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue