Merge pull request #382 from gongzhengyang/master
Rust binding supports multi cache policies and IPv6
This commit is contained in:
狮子的魂
GitHub
commit
1ac2113afc
|
|
@ -1,2 +1,3 @@
|
|||
[workspace]
|
||||
members = ["example", "xdb"]
|
||||
resolver = "2"
|
||||
members = ["example", "ip2region"]
|
||||
|
|
|
|||
|
|
@ -1,36 +1,31 @@
|
|||
# `ip2region xdb rust` 查询客户端实现
|
||||
## `ip2region xdb rust` 查询客户端实现
|
||||
|
||||
# 实现效果
|
||||
## Features
|
||||
- 支持`ip`字符串和`u32`/`u28` 数字两种类型的查询
|
||||
- 支持 IPv4 和 IPv6
|
||||
- 支持无缓存，Vector 索引缓存，全部数据缓存三种模式
|
||||
|
||||
得益于`xdb`数据存储格式设计以及`rust`编译器的高度代码优化
|
||||
## 缓存策略对比与说明
|
||||
| 缓存模式 | IPv4 数据内存占用 | IPv6 数据内存占用 | IPv4 benchmark 查询耗时 | IPv6 benchmark 查询耗时 |
|
||||
| ------------ | ----------- | ----------- | ------------------- |---------------------|
|
||||
| 无缓存 | 1-2MB | 1-2MB | 54 us | 122us |
|
||||
| vector index | 1-2MB | 1-2MB | 27 us | 100us |
|
||||
| 全部缓存 | 20 MB | 200 MB | 120 ns | 178 ns |
|
||||
|
||||
- 实现单核`CPU`下接近每秒千万级别的查询，如果是4核8线这样的`CPU`，采用`tokio`异步运行时，可以达到接近每秒4千万查询速度，查询速度取决于`CPU`物理核睿频频率
|
||||
- 达到查询稳定在`100-150ns/op`
|
||||
- `ip2region.xdb`文件会直接加载到内存，整个程序运行时候占用内存`13M`左右，即使是多线程或者异步运行时下面的高并发查询也是稳定在这个内存大小
|
||||
- 只会加载一次数据，多线程安全，可以自由使用`tokio`异步运行时或者标准库的多线程`std::thread`
|
||||
- 在 `ip2region::Searcher` 初始化的时候会产生一次 IO, 读取 `xdb` 的 header 信息以初始化 `Searcher`，header 信息主要包含了 `xdb` 的 IP 版本，该操作对后续 IP 的查询不产生性能，耗时影响，多占用约 20 Byte 的内存
|
||||
- 在无缓存模式与 `vector index` 缓存模式下，所有 `xdb` 的 IO 读取都是按需（按照 bytes offset, bytes length）读取少量信息, 都是线程安全的，可以 benchmark 测试验证
|
||||
- 在全部缓存模式下，`xdb` 文件会一次读取，加载到内存中，测试 `IPv6 xdb` 文件大约占用内存 200MB 左右，查询不频繁的话，占用内存会逐渐降低
|
||||
- 所有缓存模式下，包括初始化 `ip2region::Searcher` 过程当中，程序都是线程安全的，不存在某个全局可修改的中间变量，`ip2region::Searcher` 初始化完成以后，调用函数`search`都是使用不可变引用，同时 `ip2region::Searcher` 也可以通过 `Arc` 方式传递给不同线程使用
|
||||
|
||||
# 缓存方式说明
|
||||
|
||||
由于基于文件的查询以及缓存`VectorIndex`索引在并发较高(比如每秒上百并发)的情况下，每次查询都会从磁盘加载`ip2region.xdb`文件进入内存，由此会产生很高的磁盘`IO`以及极大的内存占用，所以决定做一次减法，不对这两种缓存进行开发，只提供缓存整个`xdb`文件的方式，以此实现最小的并发查询内存开销以及极限`CPU`性能压榨
|
||||
|
||||
# 使用方式
|
||||
## 使用方式
|
||||
|
||||
使用`cargo`新建一个项目，比如`cargo new ip-test`
|
||||
|
||||
同时把`ip2region.xdb`文件也移动到该项目根路径下，或者不移动，下面示例编译的时候注意调整`xdb_filepath`的参数值
|
||||
|
||||
配置`Cargo.toml`的`[dependencies]`如下
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
xdb = { git = "https://github.com/lionsoul2014/ip2region.git", branch = "master" }
|
||||
# 用于生成随机数
|
||||
rand = "0.8"
|
||||
# 用于初始化日志打印
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = "0.3.14"
|
||||
# 异步运行时
|
||||
tokio = { version = "1", features = ["full"]}
|
||||
ip2region = { git = "https://github.com/lionsoul2014/ip2region.git", branch = "master" }
|
||||
```
|
||||
|
||||
### 基本使用示例
|
||||
|
|
@ -38,322 +33,131 @@ tokio = { version = "1", features = ["full"]}
|
|||
编写`main.rs`
|
||||
|
||||
```rust
|
||||
use std::net::Ipv4Addr;
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use xdb::{search_by_ip, searcher_init};
|
||||
use ip2region::{CachePolicy, Searcher};
|
||||
|
||||
fn main() {
|
||||
// 配置输出日志信息
|
||||
tracing_subscriber::fmt::init();
|
||||
for cache_policy in [
|
||||
CachePolicy::NoCache,
|
||||
CachePolicy::FullMemory,
|
||||
CachePolicy::VectorIndex,
|
||||
] {
|
||||
let ipv4_seacher = Searcher::new("../ip2region/data/ip2region_v4.xdb".to_owned(), cache_policy).unwrap();
|
||||
for ip in [1_u32, 2, 3] {
|
||||
let result = ipv4_seacher.search(ip).unwrap();
|
||||
println!("CachePolicy: {cache_policy:?}, IP: {ip}, Result: {result}");
|
||||
}
|
||||
|
||||
// 初始化加载xdb文件
|
||||
let xdb_filepath = "./ip2region.xdb";
|
||||
searcher_init(Some(xdb_filepath.to_owned()));
|
||||
// 如果../data或者../../data或者../../../data下面有对应的ip2region.xdb文件
|
||||
// 初始化函数可以直接调用如下
|
||||
// searcher_init(None);
|
||||
for ip in ["1.1.1.1", "2.2.2.2"] {
|
||||
let result = ipv4_seacher.search(ip).unwrap();
|
||||
println!("CachePolicy: {cache_policy:?}, IP: {ip}, Result: {result}");
|
||||
}
|
||||
|
||||
println!("\n测试多类型查询");
|
||||
println!("{}", search_by_ip("9999999").unwrap());
|
||||
println!("{}", search_by_ip("1.0.1.0").unwrap());
|
||||
println!("{}", search_by_ip(9999999).unwrap());
|
||||
println!("{:?}", search_by_ip(Ipv4Addr::from(3333333)));
|
||||
let ipv6_seacher = Searcher::new("../ip2region/data/ip2region_v6.xdb".to_owned(), cache_policy).unwrap();
|
||||
for ip in ["2001::", "2001:4:112::"] {
|
||||
let result = ipv6_seacher.search(ip).unwrap();
|
||||
println!("CachePolicy: {cache_policy:?}, IP: {ip}, Result: {result}");
|
||||
}
|
||||
|
||||
println!("\n测试多线程初始化以及多线程查询");
|
||||
for i in 1..5 {
|
||||
thread::spawn(move || {
|
||||
// 再次初始化是没什么效果的
|
||||
searcher_init(Some(xdb_filepath.to_owned()));
|
||||
println!("in thread {i} {:?}", search_by_ip(rand::random::<u32>()));
|
||||
});
|
||||
for ip in [1_u128, 2, 3<<125] {
|
||||
let result = ipv6_seacher.search(ip).unwrap();
|
||||
println!("CachePolicy: {cache_policy:?}, IP: {ip}, Result: {result}");
|
||||
}
|
||||
}
|
||||
// 等待多线程执行结束
|
||||
thread::sleep(Duration::from_secs(1));
|
||||
|
||||
let count = 10_000_000;
|
||||
print!("\n计算千万数据总共耗时: ");
|
||||
let now = Instant::now();
|
||||
for _ in 0..count {
|
||||
search_by_ip(rand::random::<u32>()).unwrap();
|
||||
}
|
||||
println!("{:?}, ave: {:?}", now.elapsed(), now.elapsed()/count);
|
||||
|
||||
print!("\n计算千万数据，每次迭代都统计耗时的总共耗时: ");
|
||||
let mut total = Duration::from_secs(0);
|
||||
for _ in 0..count {
|
||||
let current = Instant::now();
|
||||
search_by_ip(rand::random::<u32>()).unwrap();
|
||||
total += current.elapsed();
|
||||
}
|
||||
println!("{:?}, ave: {:?}", total, total/count);
|
||||
|
||||
print!("\n计算空迭代耗时: ");
|
||||
let now = Instant::now();
|
||||
for _ in 0..count {}
|
||||
println!("{:?}", now.elapsed());
|
||||
}
|
||||
```
|
||||
|
||||
进行测试如下，需要指定`RUST_LOG`参数打印日志
|
||||
|
||||
```shell
|
||||
➜ RUST_LOG=debug cargo run -r
|
||||
Compiling ip-test v0.1.0 (/home/gong/rust-work/ip-test)
|
||||
Finished release [optimized] target(s) in 0.27s
|
||||
Running `target/release/ip-test`
|
||||
2022-12-24T03:31:22.921958Z DEBUG xdb::searcher: load xdb searcher file at ./ip2region.xdb
|
||||
|
||||
测试多类型查询
|
||||
0|0|0|内网IP|内网IP
|
||||
中国|0|福建省|福州市|电信
|
||||
0|0|0|内网IP|内网IP
|
||||
Ok("0|0|0|内网IP|内网IP")
|
||||
|
||||
测试多线程初始化以及多线程查询
|
||||
in thread 1 Ok("美国|0|新墨西哥|0|康卡斯特")
|
||||
in thread 3 Ok("0|0|0|内网IP|内网IP")
|
||||
in thread 2 Ok("土耳其|0|Ankara|0|0")
|
||||
in thread 4 Ok("爱尔兰|0|Cork|0|0")
|
||||
|
||||
计算千万数据总共耗时: 1.176241972s, ave: 117ns
|
||||
|
||||
计算千万数据，每次迭代都统计耗时的总共耗时: 1.407956755s, ave: 140ns
|
||||
|
||||
计算空迭代耗时: 40ns
|
||||
```
|
||||
|
||||
### `tokio`使用示例
|
||||
|
||||
```rust
|
||||
use std::time::Instant;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use xdb::{search_by_ip, searcher_init};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
// 配置输出debug 信息
|
||||
tracing_subscriber::fmt::init();
|
||||
searcher_init(Some("./ip2region.xdb".to_owned()));
|
||||
let main_now = Instant::now();
|
||||
let (tx, mut rx) = mpsc::channel(10);
|
||||
for i in 0..6 {
|
||||
let tx = tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let now = Instant::now();
|
||||
let count = 10_000_000;
|
||||
for _ in 0..count {
|
||||
search_by_ip(rand::random::<u32>()).unwrap();
|
||||
}
|
||||
let message = format!(
|
||||
"tokio spawn {i} over cost: {:?}, ave: {:?}",
|
||||
now.elapsed(),
|
||||
now.elapsed() / count
|
||||
);
|
||||
tx.send(message).await.unwrap();
|
||||
});
|
||||
}
|
||||
drop(tx);
|
||||
while let Some(message) = rx.recv().await {
|
||||
println!("{}", message);
|
||||
}
|
||||
println!("总共耗时: {:?}", main_now.elapsed());
|
||||
}
|
||||
```
|
||||
|
||||
开始执行测试
|
||||
|
||||
```shell
|
||||
$ RUST_LOG=debug cargo run -r
|
||||
Compiling ip-test v0.1.0 (/home/gong/rust-work/ip-test)
|
||||
Finished release [optimized] target(s) in 0.51s
|
||||
Running `target/release/ip-test`
|
||||
2022-12-24T04:05:32.876664Z DEBUG xdb::searcher: load xdb searcher file at ./ip2region.xdb
|
||||
tokio spawn 4 over cost: 1.133448675s, ave: 113ns
|
||||
tokio spawn 2 over cost: 1.133938619s, ave: 113ns
|
||||
tokio spawn 1 over cost: 1.136872027s, ave: 113ns
|
||||
tokio spawn 5 over cost: 1.173464286s, ave: 117ns
|
||||
tokio spawn 0 over cost: 1.197527014s, ave: 119ns
|
||||
tokio spawn 3 over cost: 1.26446099s, ave: 126ns
|
||||
总共耗时: 1.264631935s
|
||||
```
|
||||
|
||||
# `binding/rust`路径下面的结构说明
|
||||
|
||||
`xdb`
|
||||
|
||||
- 封装了`ip`到`region`的函数
|
||||
- 里面包含单元测试和`benchmark`测试
|
||||
|
||||
`example`
|
||||
|
||||
- 包含了命令行可执行文件生成的源码程序
|
||||
- 作为一个用于`rust`的开发集成例子
|
||||
|
||||
开始编译之后会生成如下
|
||||
|
||||
`target`
|
||||
|
||||
- 文件夹存放编译之后的文件以及编译产生的临时文件与缓存
|
||||
|
||||
`Cargo.lock`
|
||||
|
||||
- 固定`rust`第三方库的版本
|
||||
|
||||
编译生成的文件全部在`.gitignore`中有标识，不会被提交
|
||||
|
||||
# 编译程序
|
||||
|
||||
切换到 `ip2region/binding/rust` 路径，执行如下命令
|
||||
## Cache policy benchmark
|
||||
|
||||
```bash
|
||||
$ cargo build -r
|
||||
```
|
||||
|
||||
生成的二进制文件会在`./target/release/rust-example`位置
|
||||
|
||||
# 查询测试
|
||||
|
||||
切换到 `ip2region/binding/rust` 路径，执行如下命令
|
||||
|
||||
`help`输出如下
|
||||
|
||||
```shell
|
||||
$ ./target/release/rust-example query --help
|
||||
query test
|
||||
|
||||
Usage: rust-example query [OPTIONS]
|
||||
|
||||
Options:
|
||||
--db <db> 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
|
||||
-h, --help Print help information
|
||||
```
|
||||
|
||||
执行测试，使用默认`data/ip2region.txt`
|
||||
|
||||
```shell
|
||||
$ ./target/release/rust-example query --db=../../data/ip2region.xdb
|
||||
ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit
|
||||
ip2region>> 123123123
|
||||
region: Ok("美国|0|0|0|0"), took: 4.94µs
|
||||
ip2region>> 1.1.1.1
|
||||
region: Ok("澳大利亚|0|0|0|0"), took: 2.057µs
|
||||
ip2region>> 2.2.2.2
|
||||
region: Ok("法国|0|0|0|橘子电信"), took: 4.294µs
|
||||
ip2region>>
|
||||
```
|
||||
|
||||
这边发现每次查询的消耗时间都超过`1µs`，和开头所说的纳秒级查询不一致啊，这个是由于`rust`的标准库封装的`use std::time::Instant`对象是调用系统底层函数实现的，在打印过程中会存在时间误差
|
||||
|
||||
可以试着找一个新的项目
|
||||
|
||||
在`main.rs`写入如下
|
||||
|
||||
```rust
|
||||
use std::thread;
|
||||
use std::time::{Instant, Duration};
|
||||
|
||||
fn main() {
|
||||
let now = Instant::now();
|
||||
thread::sleep(Duration::from_secs(3));
|
||||
println!("{:?}", now.elapsed());
|
||||
}
|
||||
```
|
||||
|
||||
执行命令如下，发现毫秒级别是没什么问题，微秒和纳秒上面是存在误差的，详情可以参考[`rust`标准库文档的`time::Instant`章](https://rustwiki.org/zh-CN/std/time/struct.Instant.html)
|
||||
|
||||
```shell
|
||||
$ cargo run -r
|
||||
Finished release [optimized] target(s) in 0.03s
|
||||
Running `target/release/ip-test`
|
||||
3.000197389s
|
||||
```
|
||||
|
||||
# `bench`测试
|
||||
|
||||
测试平均性能
|
||||
|
||||
切换到 `ip2region/binding/rust` 路径，执行如下命令
|
||||
|
||||
`help`输出如下
|
||||
|
||||
```shell
|
||||
$ ./target/release/rust-example bench --help
|
||||
bench test
|
||||
|
||||
Usage: rust-example bench [OPTIONS] --src <src>
|
||||
|
||||
Options:
|
||||
--src <src> set this to specific source bench file
|
||||
--db <db> 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
|
||||
-h, --help Print help information
|
||||
```
|
||||
|
||||
使用默认的`ip2region`和`ip.merge.txt`
|
||||
|
||||
```shell
|
||||
$ ./target/release/rust-example bench --src=../../data/ip.merge.txt --db=../../data/ip2region.xdb
|
||||
Bench finished, total: 3419220,took: 519.820535ms ,cost: 152ns/op
|
||||
```
|
||||
|
||||
# `binding/rust`后续维护须知
|
||||
|
||||
`bingd/rust`编写了单元测试，后续开发需要保证单元测试正常
|
||||
|
||||
切换到 `ip2region/binding/rust` 路径，执行如下命令
|
||||
|
||||
```shell
|
||||
$ git lfs pull
|
||||
$ cd binding/rust/ip2region
|
||||
$ cargo test
|
||||
```
|
||||
$ cargo bench
|
||||
|
||||
需要保证查询速度不会有大幅降低，希望有朝一日，远方的朋友可以再深度优化一下，实现几十纳秒级别的查询速度
|
||||
|
||||
下面是`rust/xdb`库的第一版`benchmark`结果
|
||||
|
||||
重点关注如下
|
||||
|
||||
`search_by_ip_bench `
|
||||
|
||||
- 查询`ip`的实际调用函数`search_by_ip`
|
||||
|
||||
`get_block_by_size_bench`
|
||||
|
||||
- 获取并且计算偏移值，对应函数`get_block_by_size_bench`，和其他`binding`下的实现的`getLong`和`getShort`相似
|
||||
- 该函数会被`search_by_ip`多次调用，所以被标注为`#[inline]`使用内联优化，以此来消除函数调用产生的压栈开销
|
||||
|
||||
```shell
|
||||
$ cargo bench
|
||||
// --snip---
|
||||
search_by_ip_bench time: [116.99 ns 119.52 ns 122.31 ns]
|
||||
change: [-8.1930% -5.8295% -3.3029%] (p = 0.00 < 0.05)
|
||||
Performance has improved.
|
||||
Found 4 outliers among 100 measurements (4.00%)
|
||||
4 (4.00%) high mild
|
||||
ipv4_no_memory_bench time: [54.699 µs 57.401 µs 61.062 µs]
|
||||
Found 16 outliers among 100 measurements (16.00%)
|
||||
10 (10.00%) high mild
|
||||
6 (6.00%) high severe
|
||||
|
||||
get_block_by_size_bench time: [5.2388 ns 5.2784 ns 5.3229 ns]
|
||||
change: [-6.2649% -4.3559% -2.5539%] (p = 0.00 < 0.05)
|
||||
Performance has improved.
|
||||
Found 7 outliers among 100 measurements (7.00%)
|
||||
3 (3.00%) high mild
|
||||
4 (4.00%) high severe
|
||||
|
||||
get_full_cache_bench time: [1.4800 ns 1.5034 ns 1.5325 ns]
|
||||
change: [-17.984% -13.664% -9.1681%] (p = 0.00 < 0.05)
|
||||
Performance has improved.
|
||||
Found 15 outliers among 100 measurements (15.00%)
|
||||
3 (3.00%) high mild
|
||||
12 (12.00%) high severe
|
||||
|
||||
get_vec_index_cache_bench
|
||||
time: [1.7578 ns 1.7757 ns 1.7961 ns]
|
||||
change: [-7.5169% -4.5088% -1.7147%] (p = 0.00 < 0.05)
|
||||
Performance has improved.
|
||||
Found 5 outliers among 100 measurements (5.00%)
|
||||
3 (3.00%) high mild
|
||||
ipv4_vector_index_cache_bench
|
||||
time: [25.972 µs 26.151 µs 26.360 µs]
|
||||
Found 9 outliers among 100 measurements (9.00%)
|
||||
1 (1.00%) low severe
|
||||
6 (6.00%) high mild
|
||||
2 (2.00%) high severe
|
||||
|
||||
ipv4_full_memory_cache_bench
|
||||
time: [132.04 ns 139.48 ns 149.20 ns]
|
||||
Found 10 outliers among 100 measurements (10.00%)
|
||||
4 (4.00%) high mild
|
||||
6 (6.00%) high severe
|
||||
|
||||
ipv6_no_memory_bench time: [121.00 µs 122.14 µs 123.40 µs]
|
||||
Found 5 outliers among 100 measurements (5.00%)
|
||||
2 (2.00%) high mild
|
||||
3 (3.00%) high severe
|
||||
|
||||
ipv6_vector_index_cache_bench
|
||||
time: [96.830 µs 100.23 µs 104.81 µs]
|
||||
Found 8 outliers among 100 measurements (8.00%)
|
||||
2 (2.00%) high mild
|
||||
6 (6.00%) high severe
|
||||
|
||||
ipv6_full_memory_cache_bench
|
||||
time: [175.29 ns 178.82 ns 183.77 ns]
|
||||
Found 6 outliers among 100 measurements (6.00%)
|
||||
2 (2.00%) high mild
|
||||
4 (4.00%) high severe
|
||||
// --snip--
|
||||
```
|
||||
|
||||
## 测试与结果验证，benchmark
|
||||
```bash
|
||||
$ git lfs pull
|
||||
$ cd binding/rust/example
|
||||
$ cargo build -r
|
||||
```
|
||||
构建的执行程序位置 `binding/rust/target/release/searcher`
|
||||
|
||||
#### 测试 IPv6
|
||||
```bash
|
||||
$ cd binding/rust
|
||||
$ ./target/release/searcher --xdb='../../data/ip2region_v6.xdb' query
|
||||
|
||||
ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit
|
||||
ip2region>> 2001:5:4::
|
||||
region: Ok("荷兰|北荷兰省|阿姆斯特丹|专线用户"), took: 284.80775ms
|
||||
ip2region>> 2001::
|
||||
region: Ok("美国|加利福尼亚州|洛杉矶|专线用户"), took: 12.75µs
|
||||
ip2region>> 2001:5:6::
|
||||
region: Ok("荷兰|北荷兰省|阿姆斯特丹|专线用户"), took: 52.958µs
|
||||
ip2region>> 2001:5:5::
|
||||
region: Ok("比利时|弗拉芒大区|泽勒|专线用户"), took: 123.375µs
|
||||
ip2region>>
|
||||
```
|
||||
|
||||
#### 测试 IPv4
|
||||
```bash
|
||||
$ cd binding/rust
|
||||
$ ./target/release/searcher --xdb='../../data/ip2region_v4.xdb' query
|
||||
ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit
|
||||
ip2region>> 1.1.2.1
|
||||
region: Ok("中国|福建省|福州市|电信"), took: 5.342625ms
|
||||
ip2region>> 2.2.21.1
|
||||
region: Ok("法国|0|0|橘子电信"), took: 25.667µs
|
||||
ip2region>>
|
||||
```
|
||||
|
||||
#### Benchmark 与验证结果
|
||||
|
||||
通过 searcher 程序来测试性能，同时依据 ip sources 文件对比查询结果，检测是否存在错误
|
||||
|
||||
```bash
|
||||
$ cd binding/rust/example
|
||||
$ cargo build -r
|
||||
## 通过 data/ip2region_v4.xdb 和 data/ipv4_source.txt 进行 ipv4 的 bench 测试：
|
||||
$ RUST_LOG=debug ../target/release/searcher --xdb='../../../data/ip2region_v4.xdb' bench '../../../data/ipv4_source.txt'
|
||||
## 通过 data/ip2region_v6.xdb 和 data/ipv6_source.txt 进行 ipv6 的 bench 测试：
|
||||
$ RUST_LOG=debug ../target/release/searcher --xdb='../../../data/ip2region_v6.xdb' bench '../../../data/ipv6_source.txt'
|
||||
```
|
||||
|
|
@ -1,16 +1,16 @@
|
|||
[package]
|
||||
name = "rust-example"
|
||||
default-run = "rust-example"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
rust-version = "1.66.0"
|
||||
description = "the rust binding for ip2region"
|
||||
name = "searcher"
|
||||
default-run = "searcher"
|
||||
version = "0.2.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.89.0"
|
||||
description = "Rust binding example for ip2region"
|
||||
license = "Apache-2.0"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
xdb = { path = "../xdb" }
|
||||
clap = { version = "4.0" }
|
||||
ip2region = { path = "../ip2region" }
|
||||
clap = { version = "4.5", features = ["derive", "env"] }
|
||||
tracing-subscriber = "0.3"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = "0.3.14"
|
||||
|
|
|
|||
|
|
@ -1,27 +1,45 @@
|
|||
use clap::{Arg, ArgMatches, Command};
|
||||
use clap::{Parser, Subcommand, ValueEnum};
|
||||
|
||||
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 --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()
|
||||
/// Rust binding example for ip2region
|
||||
///
|
||||
/// e.g
|
||||
///
|
||||
/// ```
|
||||
///
|
||||
/// export XDB='../../../data/ip2region_v4.xdb' ## or export XDB='../../../data/ip2region_v6.xdb'
|
||||
///
|
||||
/// export CHECK='../../../data/ipv4_source.txt' ## or export CHECK='../../../data/ipv6_source.txt'
|
||||
///
|
||||
/// cd binding/rust/example
|
||||
///
|
||||
/// ./searcher --xdb=$XDB bench $CHECK
|
||||
///
|
||||
/// ./searcher --xdb=$XDB query
|
||||
///
|
||||
/// ```
|
||||
#[derive(Parser)]
|
||||
pub struct Command {
|
||||
/// xdb filepath, e.g. `../../../data/ip2region_v4.xdb` or `../../../data/ip2region_v6.xdb`
|
||||
#[arg(long, env = "XDB")]
|
||||
pub xdb: String,
|
||||
#[arg(long, value_enum, default_value_t = CmdCachePolicy::FullMemory)]
|
||||
pub cache_policy: CmdCachePolicy,
|
||||
#[clap(subcommand)]
|
||||
pub action: Action,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum Action {
|
||||
/// Bench the ip search and output performance info
|
||||
Bench { check_file: String },
|
||||
/// Interactive input and output, querying one IP and get result at a time
|
||||
Query,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, ValueEnum, Clone, Copy, Default)]
|
||||
pub enum CmdCachePolicy {
|
||||
#[default]
|
||||
FullMemory,
|
||||
NoCache,
|
||||
VectorIndex,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,73 +1,82 @@
|
|||
extern crate core;
|
||||
|
||||
use std::fs::File;
|
||||
use std::io::Read;
|
||||
use std::io::Write;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::net::IpAddr;
|
||||
use std::str::FromStr;
|
||||
use std::time::Instant;
|
||||
|
||||
use clap::ArgMatches;
|
||||
use clap::Parser;
|
||||
use ip2region::{CachePolicy, Searcher};
|
||||
use tracing::info;
|
||||
|
||||
use xdb::{search_by_ip, searcher_init};
|
||||
use crate::cmd::{Action, CmdCachePolicy, Command};
|
||||
|
||||
mod cmd;
|
||||
|
||||
/// 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();
|
||||
macro_rules! perform_check {
|
||||
($searcher:expr, $start_ip:expr, $end_ip:expr, $check:expr) => {{
|
||||
let start_ip = $start_ip;
|
||||
let end_ip = $end_ip;
|
||||
|
||||
let mid_ip = (start_ip >> 1) + (end_ip >> 1);
|
||||
|
||||
let mut checked = 0;
|
||||
let checks = [
|
||||
start_ip,
|
||||
(start_ip >> 1) + (mid_ip >> 1),
|
||||
mid_ip,
|
||||
(mid_ip >> 1) + (end_ip >> 1),
|
||||
end_ip,
|
||||
];
|
||||
for ip in checks.iter() {
|
||||
if *ip < start_ip || *ip > end_ip {
|
||||
// IP not in start - end ip
|
||||
// This happens when start ip equals end ip
|
||||
continue;
|
||||
}
|
||||
let result = $searcher.search(*ip).unwrap();
|
||||
assert_eq!(result.as_str(), $check);
|
||||
checked += 1;
|
||||
}
|
||||
checked
|
||||
}};
|
||||
}
|
||||
|
||||
fn bench_test(src_filepath: &str) {
|
||||
fn check(searcher: &Searcher, start_ip: IpAddr, end_ip: IpAddr, check: &str) -> usize {
|
||||
match (start_ip, end_ip) {
|
||||
(IpAddr::V4(original_start_ip), IpAddr::V4(original_end_ip)) => {
|
||||
let start_ip = u32::from(original_start_ip);
|
||||
let end_ip = u32::from(original_end_ip);
|
||||
perform_check!(searcher, start_ip, end_ip, check)
|
||||
}
|
||||
(IpAddr::V6(original_start_ip), IpAddr::V6(original_end_ip)) => {
|
||||
let start_ip = u128::from(original_start_ip);
|
||||
let end_ip = u128::from(original_end_ip);
|
||||
perform_check!(searcher, start_ip, end_ip, check)
|
||||
}
|
||||
_ => panic!("invalid start ip and end ip"),
|
||||
}
|
||||
}
|
||||
|
||||
fn bench(searcher: &Searcher, check_filepath: &str) {
|
||||
let file = File::open(check_filepath).unwrap();
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
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;
|
||||
}
|
||||
for line in reader.lines().map_while(Result::ok) {
|
||||
let ip_test_line = line.splitn(3, '|').collect::<Vec<&str>>();
|
||||
if ip_test_line.len() != 3 {
|
||||
panic!("this line {line} don`t have enough `|` for spilt");
|
||||
}
|
||||
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,
|
||||
] {
|
||||
let result = search_by_ip(ip).unwrap();
|
||||
assert_eq!(result.as_str(), ip_test_line[2]);
|
||||
count += 1;
|
||||
if ip_test_line.len() == 3 {
|
||||
let start_ip = IpAddr::from_str(ip_test_line[0]).unwrap();
|
||||
let end_ip = IpAddr::from_str(ip_test_line[1]).unwrap();
|
||||
count += check(searcher, start_ip, end_ip, ip_test_line[2]);
|
||||
}
|
||||
}
|
||||
println!(
|
||||
"Bench finished, total: {count},\
|
||||
took: {:?} ,\
|
||||
cost: {:?}/op",
|
||||
now.elapsed(),
|
||||
now.elapsed() / count
|
||||
)
|
||||
info!(count, took=?now.elapsed(), avg_took=?(now.elapsed() / (count as u32)), "Benchmark finished");
|
||||
}
|
||||
|
||||
fn query_test() {
|
||||
fn query(searcher: &Searcher) {
|
||||
println!("ip2region xdb searcher test program, type `quit` or `Ctrl + c` to exit");
|
||||
loop {
|
||||
print!("ip2region>> ");
|
||||
|
|
@ -79,32 +88,25 @@ fn query_test() {
|
|||
}
|
||||
let line = line.trim();
|
||||
let now = Instant::now();
|
||||
let result = search_by_ip(line);
|
||||
let result = searcher.search(line);
|
||||
let cost = now.elapsed();
|
||||
println!("region: {result:?}, took: {cost:?}",);
|
||||
}
|
||||
}
|
||||
|
||||
fn matches_for_searcher(matches: &ArgMatches) {
|
||||
if let Some(xdb_filepath) = matches.get_one::<String>("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::<String>("src").unwrap();
|
||||
tracing_subscriber::fmt::init();
|
||||
|
||||
bench_test(src_filepath);
|
||||
}
|
||||
let cmd = Command::parse();
|
||||
let cache_policy = match cmd.cache_policy {
|
||||
CmdCachePolicy::FullMemory => CachePolicy::FullMemory,
|
||||
CmdCachePolicy::VectorIndex => CachePolicy::VectorIndex,
|
||||
CmdCachePolicy::NoCache => CachePolicy::NoCache,
|
||||
};
|
||||
|
||||
if let Some(sub_matches) = matches.subcommand_matches("query") {
|
||||
matches_for_searcher(sub_matches);
|
||||
query_test()
|
||||
let searcher = Searcher::new(cmd.xdb, cache_policy).unwrap();
|
||||
match cmd.action {
|
||||
Action::Bench { check_file } => bench(&searcher, &check_file),
|
||||
Action::Query => query(&searcher),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,21 +1,22 @@
|
|||
[package]
|
||||
name = "xdb"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
rust-version = "1.66.0"
|
||||
description = "the rust binding for ip2region"
|
||||
name = "ip2region"
|
||||
version = "0.2.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.89.0"
|
||||
description = "The rust binding for ip2region"
|
||||
license = "Apache-2.0"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
once_cell = "1.16"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = "0.3.14"
|
||||
thiserror = "2"
|
||||
num-traits = "0.2"
|
||||
num-derive = "0.4"
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = "0.4"
|
||||
rand = "0.8"
|
||||
criterion = "0.7"
|
||||
rand = "0.9"
|
||||
|
||||
[[bench]]
|
||||
name = "search"
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
use std::net::Ipv6Addr;
|
||||
use std::ops::Range;
|
||||
use std::str::FromStr;
|
||||
|
||||
use criterion::{Criterion, criterion_group, criterion_main};
|
||||
|
||||
use ip2region::{CachePolicy, Searcher};
|
||||
|
||||
macro_rules! bench_search {
|
||||
($name:ident, $xdb:expr, $cache_policy:expr, $range:ident) => {
|
||||
fn $name(c: &mut Criterion) {
|
||||
let searcher = Searcher::new($xdb.to_owned(), $cache_policy).unwrap();
|
||||
let range = $range();
|
||||
|
||||
c.bench_function(stringify!($name), |b| {
|
||||
b.iter(|| {
|
||||
searcher.search(rand::random_range(range.clone())).unwrap();
|
||||
})
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
fn ipv4_range() -> Range<u32> {
|
||||
0..((1_u64 << 32) - 1) as u32
|
||||
}
|
||||
|
||||
/// The range of IPv6 is too large, and the value range needs to be limited to
|
||||
/// make the benchmark test results closer to the production environment
|
||||
fn ipv6_range() -> Range<u128> {
|
||||
let start = u128::from(Ipv6Addr::from_str("2000::").unwrap());
|
||||
let end = u128::from(Ipv6Addr::from_str("2004::").unwrap());
|
||||
start..end
|
||||
}
|
||||
|
||||
const IPV4_XDB: &str = "../../../data/ip2region_v4.xdb";
|
||||
const IPV6_XDB: &str = "../../../data/ip2region_v6.xdb";
|
||||
|
||||
bench_search!(
|
||||
ipv4_no_memory_bench,
|
||||
IPV4_XDB,
|
||||
CachePolicy::NoCache,
|
||||
ipv4_range
|
||||
);
|
||||
bench_search!(
|
||||
ipv4_vector_index_cache_bench,
|
||||
IPV4_XDB,
|
||||
CachePolicy::VectorIndex,
|
||||
ipv4_range
|
||||
);
|
||||
bench_search!(
|
||||
ipv4_full_memory_cache_bench,
|
||||
IPV4_XDB,
|
||||
CachePolicy::FullMemory,
|
||||
ipv4_range
|
||||
);
|
||||
bench_search!(
|
||||
ipv6_no_memory_bench,
|
||||
IPV6_XDB,
|
||||
CachePolicy::NoCache,
|
||||
ipv6_range
|
||||
);
|
||||
bench_search!(
|
||||
ipv6_vector_index_cache_bench,
|
||||
IPV6_XDB,
|
||||
CachePolicy::VectorIndex,
|
||||
ipv6_range
|
||||
);
|
||||
bench_search!(
|
||||
ipv6_full_memory_cache_bench,
|
||||
IPV6_XDB,
|
||||
CachePolicy::FullMemory,
|
||||
ipv6_range
|
||||
);
|
||||
|
||||
criterion_group!(
|
||||
benches,
|
||||
ipv4_no_memory_bench,
|
||||
ipv4_vector_index_cache_bench,
|
||||
ipv4_full_memory_cache_bench,
|
||||
ipv6_no_memory_bench,
|
||||
ipv6_vector_index_cache_bench,
|
||||
ipv6_full_memory_cache_bench
|
||||
);
|
||||
criterion_main!(benches);
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Ip2RegionError {
|
||||
#[error("Io error: {0}")]
|
||||
IoError(#[from] std::io::Error),
|
||||
|
||||
#[error("From UTF-8 error: {0}")]
|
||||
Utf8Error(#[from] std::string::FromUtf8Error),
|
||||
|
||||
#[error("Parse invalid IP address")]
|
||||
ParseIpaddressFailed,
|
||||
|
||||
#[error("No matched Ipaddress")]
|
||||
NoMatchedIP,
|
||||
|
||||
#[error("Header parse error: {0}")]
|
||||
HeaderParsed(String),
|
||||
|
||||
#[error("Searcher load IPv4 data, couldn't search IPv6 data")]
|
||||
OnlyIPv4Version,
|
||||
|
||||
#[error("Searcher load IPv6 data, couldn't search IPv4 data")]
|
||||
OnlyIPv6Version,
|
||||
|
||||
#[error("Try from slice failed")]
|
||||
TryFromSliceFailed(#[from] std::array::TryFromSliceError),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Ip2RegionError>;
|
||||
|
|
@ -0,0 +1,86 @@
|
|||
use num_derive::FromPrimitive;
|
||||
use num_traits::FromPrimitive;
|
||||
|
||||
use crate::error::Ip2RegionError;
|
||||
|
||||
pub const HEADER_INFO_LENGTH: usize = 256;
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug)]
|
||||
pub struct Header {
|
||||
version: u16,
|
||||
index_policy: IndexPolicy,
|
||||
create_time: u32,
|
||||
start_index_ptr: u32,
|
||||
end_index_ptr: u32,
|
||||
ip_version: IpVersion,
|
||||
runtime_ptr_bytes: u16,
|
||||
}
|
||||
|
||||
impl TryFrom<&[u8; 256]> for Header {
|
||||
type Error = Ip2RegionError;
|
||||
|
||||
fn try_from(value: &[u8; 256]) -> Result<Self, Self::Error> {
|
||||
if value.len() < 20 {
|
||||
return Err(Ip2RegionError::HeaderParsed(
|
||||
"Header bytes too short".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let index_policy_value = u16::from_le_bytes([value[2], value[3]]);
|
||||
let ip_version_value = u16::from_le_bytes([value[16], value[17]]);
|
||||
|
||||
Ok(Header {
|
||||
version: u16::from_le_bytes([value[0], value[1]]),
|
||||
index_policy: IndexPolicy::from_u16(index_policy_value).ok_or_else(|| {
|
||||
Ip2RegionError::HeaderParsed(format!(
|
||||
"Header index policy invalid: {index_policy_value}"
|
||||
))
|
||||
})?,
|
||||
create_time: u32::from_le_bytes([value[4], value[5], value[6], value[7]]),
|
||||
start_index_ptr: u32::from_le_bytes([value[8], value[9], value[10], value[11]]),
|
||||
end_index_ptr: u32::from_le_bytes([value[12], value[13], value[14], value[15]]),
|
||||
|
||||
ip_version: IpVersion::from_u16(ip_version_value).ok_or_else(|| {
|
||||
Ip2RegionError::HeaderParsed(format!(
|
||||
"Header ip version invalid: {ip_version_value}"
|
||||
))
|
||||
})?,
|
||||
runtime_ptr_bytes: u16::from_le_bytes([value[18], value[19]]),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(FromPrimitive, Debug)]
|
||||
#[repr(u16)]
|
||||
pub enum IndexPolicy {
|
||||
VectorIndex = 1,
|
||||
BTreeIndex = 2,
|
||||
}
|
||||
|
||||
#[derive(FromPrimitive, Debug)]
|
||||
#[repr(u16)]
|
||||
pub enum IpVersion {
|
||||
V4 = 4,
|
||||
V6 = 6,
|
||||
}
|
||||
|
||||
impl Header {
|
||||
pub fn ip_bytes_len(&self) -> usize {
|
||||
match &self.ip_version {
|
||||
IpVersion::V4 => 4,
|
||||
IpVersion::V6 => 16,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn segment_index_size(&self) -> usize {
|
||||
match &self.ip_version {
|
||||
IpVersion::V4 => 14,
|
||||
IpVersion::V6 => 38,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ip_version(&self) -> &IpVersion {
|
||||
&self.ip_version
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
use std::borrow::Cow;
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
|
||||
use std::str::FromStr;
|
||||
|
||||
use crate::error::{Ip2RegionError, Result};
|
||||
|
||||
pub trait IpValueExt {
|
||||
fn to_ipaddr(self) -> Result<IpAddr>;
|
||||
}
|
||||
|
||||
impl IpValueExt for &str {
|
||||
fn to_ipaddr(self) -> Result<IpAddr> {
|
||||
IpAddr::from_str(self).map_err(|_| Ip2RegionError::ParseIpaddressFailed)
|
||||
}
|
||||
}
|
||||
|
||||
impl IpValueExt for u32 {
|
||||
fn to_ipaddr(self) -> Result<IpAddr> {
|
||||
Ok(IpAddr::V4(Ipv4Addr::from(self)))
|
||||
}
|
||||
}
|
||||
|
||||
impl IpValueExt for Ipv4Addr {
|
||||
fn to_ipaddr(self) -> Result<IpAddr> {
|
||||
Ok(IpAddr::V4(self))
|
||||
}
|
||||
}
|
||||
|
||||
impl IpValueExt for Ipv6Addr {
|
||||
fn to_ipaddr(self) -> Result<IpAddr> {
|
||||
Ok(IpAddr::V6(self))
|
||||
}
|
||||
}
|
||||
|
||||
impl IpValueExt for u128 {
|
||||
fn to_ipaddr(self) -> Result<IpAddr> {
|
||||
Ok(IpAddr::V6(Ipv6Addr::from(self)))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CompareExt {
|
||||
fn ip_lt(&self, other: Cow<'_, [u8]>) -> bool;
|
||||
fn ip_gt(&self, other: Cow<'_, [u8]>) -> bool;
|
||||
}
|
||||
|
||||
impl CompareExt for IpAddr {
|
||||
fn ip_lt(&self, other: Cow<'_, [u8]>) -> bool {
|
||||
match self {
|
||||
IpAddr::V4(ip) => ip.octets() < [other[3], other[2], other[1], other[0]],
|
||||
IpAddr::V6(ip) => ip.octets() < other[0..16].try_into().unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ip_gt(&self, other: Cow<'_, [u8]>) -> bool {
|
||||
match self {
|
||||
IpAddr::V4(ip) => ip.octets() > [other[3], other[2], other[1], other[0]],
|
||||
IpAddr::V6(ip) => ip.octets() > other[0..16].try_into().unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
mod error;
|
||||
mod header;
|
||||
mod ip_value;
|
||||
mod searcher;
|
||||
|
||||
pub use searcher::{CachePolicy, Searcher};
|
||||
pub use ip_value::IpValueExt;
|
||||
|
|
@ -0,0 +1,243 @@
|
|||
use std::borrow::Cow;
|
||||
use std::fmt::Display;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::net::IpAddr;
|
||||
use std::path::Path;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use tracing::{debug, trace, warn};
|
||||
|
||||
use crate::error::{Ip2RegionError, Result};
|
||||
use crate::header::{HEADER_INFO_LENGTH, Header, IpVersion};
|
||||
use crate::ip_value::{CompareExt, IpValueExt};
|
||||
|
||||
const VECTOR_INDEX_LENGTH: usize = 256 * 256 * 8;
|
||||
const VECTOR_INDEX_COLS: usize = 256;
|
||||
const VECTOR_INDEX_SIZE: usize = 8;
|
||||
|
||||
pub struct Searcher {
|
||||
pub filepath: String,
|
||||
pub cache_policy: CachePolicy,
|
||||
pub header: Header,
|
||||
vector_cache: OnceLock<Vec<u8>>,
|
||||
full_cache: OnceLock<Vec<u8>>,
|
||||
}
|
||||
|
||||
#[derive(PartialEq, Debug, Copy, Clone)]
|
||||
pub enum CachePolicy {
|
||||
NoCache,
|
||||
VectorIndex,
|
||||
FullMemory,
|
||||
}
|
||||
|
||||
impl Searcher {
|
||||
pub fn new(filepath: String, cache_policy: CachePolicy) -> Result<Self> {
|
||||
let mut file = File::open(Path::new(&filepath))?;
|
||||
let mut buf = [0; HEADER_INFO_LENGTH];
|
||||
file.read_exact(&mut buf)?;
|
||||
|
||||
let header = Header::try_from(&buf)?;
|
||||
debug!(?header, "Load xdb file with header");
|
||||
|
||||
Ok(Self {
|
||||
filepath,
|
||||
cache_policy,
|
||||
header,
|
||||
vector_cache: OnceLock::new(),
|
||||
full_cache: OnceLock::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn search<T>(&self, ip: T) -> Result<String>
|
||||
where
|
||||
T: IpValueExt + Display,
|
||||
{
|
||||
let ip = ip.to_ipaddr()?;
|
||||
|
||||
let (il0, il1) = match (ip, self.header.ip_version()) {
|
||||
(IpAddr::V6(ip), IpVersion::V6) => (ip.octets()[0], ip.octets()[1]),
|
||||
(IpAddr::V4(ip), IpVersion::V4) => (ip.octets()[0], ip.octets()[1]),
|
||||
(_, IpVersion::V4) => return Err(Ip2RegionError::OnlyIPv4Version),
|
||||
(_, IpVersion::V6) => return Err(Ip2RegionError::OnlyIPv6Version),
|
||||
};
|
||||
|
||||
let start_point = VECTOR_INDEX_SIZE * ((il0 as usize) * VECTOR_INDEX_COLS + (il1 as usize));
|
||||
let vector_index = self.vector_index()?;
|
||||
let start_ptr =
|
||||
u32::from_le_bytes(vector_index[start_point..start_point + 4].try_into()?) as usize;
|
||||
let end_ptr =
|
||||
u32::from_le_bytes(vector_index[start_point + 4..start_point + 8].try_into()?) as usize;
|
||||
|
||||
// Binary search the segment index to get the region
|
||||
let segment_index_size = self.header.segment_index_size();
|
||||
let ip_bytes_len = self.header.ip_bytes_len();
|
||||
let ip_end_offset = ip_bytes_len * 2;
|
||||
|
||||
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 = self.read_buf(offset, segment_index_size)?;
|
||||
if ip.ip_lt(Cow::Borrowed(&buffer_ip_value[0..ip_bytes_len])) {
|
||||
right = mid - 1;
|
||||
} else if ip.ip_gt(Cow::Borrowed(&buffer_ip_value[ip_bytes_len..ip_end_offset])) {
|
||||
left = mid + 1;
|
||||
} else {
|
||||
let data_length = u16::from_le_bytes([
|
||||
buffer_ip_value[ip_end_offset],
|
||||
buffer_ip_value[ip_end_offset + 1],
|
||||
]);
|
||||
let data_offset = u32::from_le_bytes(
|
||||
buffer_ip_value[ip_end_offset + 2..ip_end_offset + 6].try_into()?,
|
||||
);
|
||||
let result = String::from_utf8(
|
||||
self.read_buf(data_offset as usize, data_length as usize)?
|
||||
.to_vec(),
|
||||
)?;
|
||||
return Ok(result);
|
||||
}
|
||||
}
|
||||
Err(Ip2RegionError::NoMatchedIP)
|
||||
}
|
||||
|
||||
pub fn vector_index(&self) -> Result<Cow<'_, [u8]>> {
|
||||
if self.cache_policy.eq(&CachePolicy::NoCache) {
|
||||
return self.read_buf(HEADER_INFO_LENGTH, VECTOR_INDEX_LENGTH);
|
||||
}
|
||||
|
||||
match self.vector_cache.get() {
|
||||
None => {
|
||||
debug!("Load vector index cache");
|
||||
let data = self
|
||||
.read_buf(HEADER_INFO_LENGTH, VECTOR_INDEX_LENGTH)?
|
||||
.to_vec();
|
||||
let _ = self
|
||||
.vector_cache
|
||||
.set(data)
|
||||
.inspect_err(|_| warn!("Vector index cache already initialized"));
|
||||
|
||||
// Safety: vector cache checked and set for empty before
|
||||
let cache = self.vector_cache.get().unwrap();
|
||||
Ok(Cow::Borrowed(cache))
|
||||
}
|
||||
Some(cache) => Ok(Cow::Borrowed(cache)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_buf(&self, offset: usize, size: usize) -> Result<Cow<'_, [u8]>> {
|
||||
trace!(offset, size = size, "Read buffer");
|
||||
if self.cache_policy.ne(&CachePolicy::FullMemory) {
|
||||
debug!(filepath=?self.filepath, offset=offset, size=size, "Read buf without cache");
|
||||
let mut file = File::open(&self.filepath)?;
|
||||
file.seek(SeekFrom::Start(offset as u64))?;
|
||||
|
||||
let mut buf = vec![0u8; size];
|
||||
file.take(size as u64).read_exact(&mut buf)?;
|
||||
return Ok(Cow::from(buf));
|
||||
}
|
||||
|
||||
match self.full_cache.get() {
|
||||
None => {
|
||||
debug!(filepath=?self.filepath, "Load full cache");
|
||||
let mut file = File::open(&self.filepath)?;
|
||||
let mut buf = Vec::new();
|
||||
file.read_to_end(&mut buf)?;
|
||||
let _ = self
|
||||
.full_cache
|
||||
.set(buf)
|
||||
.inspect_err(|_| warn!("Full cache already initialized"));
|
||||
|
||||
// Safety: FULL_CACHE checked and set for empty before
|
||||
let cache = self.full_cache.get().unwrap();
|
||||
Ok(Cow::from(&cache[offset..offset + size]))
|
||||
}
|
||||
Some(cache) => {
|
||||
let data = Cow::from(&cache[offset..offset + size]);
|
||||
Ok(data)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::fs::File;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::str::FromStr;
|
||||
|
||||
use super::*;
|
||||
|
||||
// Test ipv6 need after run command `git lfs pull`
|
||||
const IPV4_XDB_PATH: &str = "../../../data/ip2region_v4.xdb";
|
||||
const IPV4_CHECK_PATH: &str = "../../../data/ipv4_source.txt";
|
||||
const IPV6_XDB_PATH: &str = "../../../data/ip2region_v6.xdb";
|
||||
const IPV6_CHECK_PATH: &str = "../../../data/ipv6_source.txt";
|
||||
|
||||
///test all types find correct
|
||||
#[test]
|
||||
fn test_multi_type_ip() {
|
||||
for cache_policy in [
|
||||
CachePolicy::NoCache,
|
||||
CachePolicy::FullMemory,
|
||||
CachePolicy::VectorIndex,
|
||||
] {
|
||||
let searcher = Searcher::new(IPV4_XDB_PATH.to_owned(), cache_policy).unwrap();
|
||||
searcher.search("1.0.1.0").unwrap();
|
||||
searcher.search("1.0.1.2").unwrap();
|
||||
searcher.search(0u32).unwrap();
|
||||
|
||||
let searcher = Searcher::new(IPV6_XDB_PATH.to_owned(), cache_policy).unwrap();
|
||||
searcher.search("2c0f:fff1::").unwrap();
|
||||
searcher.search("2c0f:fff1::1").unwrap();
|
||||
searcher.search(111u128).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn match_ip_correct(xdb_filepath: &str, check_path: &str, cache_policy: CachePolicy) {
|
||||
let searcher = Searcher::new(xdb_filepath.to_owned(), cache_policy).unwrap();
|
||||
|
||||
let file = File::open(check_path).unwrap();
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
for line in reader.lines().take(10_000) {
|
||||
let line = line.unwrap();
|
||||
|
||||
if !line.contains("|") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let ip_test_line = line.splitn(3, "|").collect::<Vec<&str>>();
|
||||
let start_ip = IpAddr::from_str(ip_test_line[0]).unwrap();
|
||||
let end_ip = IpAddr::from_str(ip_test_line[1]).unwrap();
|
||||
for _ in 0..3 {
|
||||
let result = match (start_ip, end_ip) {
|
||||
(IpAddr::V4(start), IpAddr::V4(end)) => {
|
||||
let value = rand::random_range(u32::from(start)..u32::from(end) + 1);
|
||||
searcher.search(value).unwrap()
|
||||
}
|
||||
(IpAddr::V6(start), IpAddr::V6(end)) => {
|
||||
let value = rand::random_range(u128::from(start)..u128::from(end) + 1);
|
||||
searcher.search(value).unwrap()
|
||||
}
|
||||
_ => panic!("invalid ip address"),
|
||||
};
|
||||
assert_eq!(result.as_str(), ip_test_line[2])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_match_ip_correct() {
|
||||
for cache_policy in [
|
||||
CachePolicy::NoCache,
|
||||
CachePolicy::FullMemory,
|
||||
CachePolicy::VectorIndex,
|
||||
] {
|
||||
match_ip_correct(IPV4_XDB_PATH, IPV4_CHECK_PATH, cache_policy);
|
||||
match_ip_correct(IPV6_XDB_PATH, IPV6_CHECK_PATH, cache_policy);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,52 +0,0 @@
|
|||
use criterion::{black_box, criterion_group, criterion_main, Criterion};
|
||||
use rand;
|
||||
|
||||
use xdb::searcher::{
|
||||
get_block_by_size, get_full_cache, get_vector_index_cache, search_by_ip, searcher_init,
|
||||
};
|
||||
|
||||
fn search_by_ip_bench(c: &mut Criterion) {
|
||||
c.bench_function("search_by_ip_bench", |b| {
|
||||
searcher_init(None);
|
||||
b.iter(|| {
|
||||
search_by_ip(rand::random::<u32>()).unwrap();
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
fn get_block_by_size_bench(c: &mut Criterion) {
|
||||
c.bench_function("get_block_by_size_bench", |b| {
|
||||
b.iter(|| {
|
||||
black_box(get_block_by_size(
|
||||
get_full_cache(),
|
||||
rand::random::<u16>() as usize,
|
||||
4,
|
||||
));
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
fn get_full_cache_bench(c: &mut Criterion) {
|
||||
c.bench_function("get_full_cache_bench", |b| {
|
||||
b.iter(|| {
|
||||
black_box(get_full_cache());
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
fn get_vec_index_cache_bench(c: &mut Criterion) {
|
||||
c.bench_function("get_vec_index_cache_bench", |b| {
|
||||
b.iter(|| {
|
||||
black_box(get_vector_index_cache());
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
criterion_group!(
|
||||
benches,
|
||||
search_by_ip_bench,
|
||||
get_block_by_size_bench,
|
||||
get_full_cache_bench,
|
||||
get_vec_index_cache_bench,
|
||||
);
|
||||
criterion_main!(benches);
|
||||
|
|
@ -1,61 +0,0 @@
|
|||
use std::error::Error;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::str::FromStr;
|
||||
|
||||
pub trait ToUIntIP {
|
||||
fn to_u32_ip(&self) -> Result<u32, Box<dyn Error>>;
|
||||
}
|
||||
|
||||
impl ToUIntIP for u32 {
|
||||
fn to_u32_ip(&self) -> Result<u32, Box<dyn Error>> {
|
||||
Ok(self.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
impl ToUIntIP for &str {
|
||||
fn to_u32_ip(&self) -> Result<u32, Box<dyn Error>> {
|
||||
if let Ok(ip_addr) = Ipv4Addr::from_str(self) {
|
||||
return Ok(u32::from(ip_addr));
|
||||
}
|
||||
Ok(self.parse::<u32>()?)
|
||||
}
|
||||
}
|
||||
|
||||
impl ToUIntIP for Ipv4Addr {
|
||||
fn to_u32_ip(&self) -> Result<u32, Box<dyn Error>> {
|
||||
Ok(u32::from(*self))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test_ip {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_ip_str_2_u32() {
|
||||
let ip_str = "1.1.1.1";
|
||||
let result = ip_str.to_u32_ip().unwrap();
|
||||
assert_eq!(result, 1 << 24 | 1 << 16 | 1 << 8 | 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ip_u32_str() {
|
||||
let ip = "12";
|
||||
let result = ip.to_u32_ip().unwrap();
|
||||
assert_eq!(result, 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ip_u32() {
|
||||
let ip: u32 = 33;
|
||||
let result = ip.to_u32_ip().unwrap();
|
||||
assert_eq!(result, 33);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ip_addr() {
|
||||
let ip = Ipv4Addr::from_str("0.0.3.12").unwrap();
|
||||
let result = ip.to_u32_ip().unwrap();
|
||||
assert_eq!(result, 3 << 8 | 12)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +0,0 @@
|
|||
mod ip_value;
|
||||
pub use self::ip_value::ToUIntIP;
|
||||
pub mod searcher;
|
||||
pub use searcher::{search_by_ip, searcher_init};
|
||||
|
|
@ -1,168 +0,0 @@
|
|||
use std::error::Error;
|
||||
use std::fmt::Display;
|
||||
use std::fs::File;
|
||||
use std::io::Read;
|
||||
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";
|
||||
|
||||
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,
|
||||
{
|
||||
let ip = ip.to_u32_ip()?;
|
||||
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);
|
||||
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())
|
||||
}
|
||||
|
||||
/// it will check ../data/ip2region.xdb, ../../data/ip2region.xdb, ../../../data/ip2region.xdb
|
||||
fn default_detect_xdb_file() -> Result<String, Box<dyn Error>> {
|
||||
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) << (index << 3);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
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 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_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();
|
||||
f.read_to_end(&mut buffer).expect("load file error");
|
||||
buffer
|
||||
}
|
||||
|
||||
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 super::*;
|
||||
|
||||
///test all types find correct
|
||||
#[test]
|
||||
fn test_multi_type_ip() {
|
||||
searcher_init(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();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_match_all_ip_correct() {
|
||||
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();
|
||||
for line in contents.split("\n") {
|
||||
if !line.contains("|") {
|
||||
continue;
|
||||
}
|
||||
let ip_test_line = line.splitn(3, "|").collect::<Vec<&str>>();
|
||||
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() {
|
||||
searcher_init(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();
|
||||
}
|
||||
|
||||
#[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();
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue