feat(erlang): make worker xdb version-aware for IPv4/IPv6
This commit is contained in:
Alice39s
parent
2f0dc5f33f
commit
241831a6f0
|
|
@ -42,9 +42,23 @@ init([]) ->
|
|||
%%
|
||||
create_table() ->
|
||||
Opts = [named_table, set, public, {read_concurrency, true}, {keypos, 1}],
|
||||
ets:new(?XDB_VECTOR_INDEX, Opts),
|
||||
ets:new(?XDB_SEGMENT_INDEX, Opts),
|
||||
ets:new(?IP2REGION_CACHE, Opts).
|
||||
%% Legacy tables (kept for backward compatibility)
|
||||
ensure_table(?XDB_VECTOR_INDEX, Opts),
|
||||
ensure_table(?XDB_SEGMENT_INDEX, Opts),
|
||||
ensure_table(?IP2REGION_CACHE, Opts),
|
||||
%% Version-specific tables for dual-stack support
|
||||
ensure_table(?XDB_VECTOR_INDEX_V4, Opts),
|
||||
ensure_table(?XDB_VECTOR_INDEX_V6, Opts),
|
||||
ensure_table(?XDB_SEGMENT_INDEX_V4, Opts),
|
||||
ensure_table(?XDB_SEGMENT_INDEX_V6, Opts),
|
||||
ensure_table(?IP2REGION_CACHE_V4, Opts),
|
||||
ensure_table(?IP2REGION_CACHE_V6, Opts).
|
||||
|
||||
ensure_table(Name, Opts) ->
|
||||
case ets:whereis(Name) of
|
||||
undefined -> ets:new(Name, Opts);
|
||||
_ -> ok
|
||||
end.
|
||||
|
||||
start_ip2region_pool(Sup) ->
|
||||
{ok, PoolArgsCfg} = application:get_env(poolargs),
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
%% license that can be found in the LICENSE file.
|
||||
%%
|
||||
%% @doc
|
||||
%% ip2region xdb client worker
|
||||
%% ip2region xdb client worker, now version-aware (IPv4/IPv6).
|
||||
%% @end
|
||||
%%%-------------------------------------------------------------------
|
||||
-module(ip2region_worker).
|
||||
|
|
@ -15,7 +15,12 @@
|
|||
-export([start/1, stop/1, start_link/1]).
|
||||
-export([search/2]).
|
||||
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
|
||||
-record(state, {xdb_fd}).
|
||||
|
||||
-record(state, {
|
||||
xdb_fd,
|
||||
version :: ipv4 | ipv6,
|
||||
segment_index_size :: pos_integer()
|
||||
}).
|
||||
|
||||
%%==========================================
|
||||
%% API
|
||||
|
|
@ -28,7 +33,6 @@ start_link(Args) ->
|
|||
Opts = [{spawn_opt, [{min_heap_size, 6000}]}],
|
||||
gen_server:start_link(?MODULE, Args, Opts).
|
||||
|
||||
|
||||
stop(Pid) ->
|
||||
gen_server:call(Pid, stop).
|
||||
|
||||
|
|
@ -38,7 +42,7 @@ search(Pid, Ip) ->
|
|||
%%==========================================
|
||||
%% gen_server callbacks
|
||||
%% =========================================
|
||||
init(_Args) ->
|
||||
init(Args) ->
|
||||
process_flag(trap_exit, true),
|
||||
AppName =
|
||||
case application:get_application() of
|
||||
|
|
@ -46,11 +50,27 @@ init(_Args) ->
|
|||
_ -> ?APP_NAME
|
||||
end,
|
||||
PrivDir = code:priv_dir(AppName),
|
||||
XdbFileName = filename:join([PrivDir, "ip2region.xdb"]),
|
||||
XdbFileName =
|
||||
case proplists:get_value(xdb_file, Args) of
|
||||
undefined -> filename:join([PrivDir, "ip2region.xdb"]);
|
||||
Path ->
|
||||
case filename:pathtype(Path) of
|
||||
absolute -> Path;
|
||||
_ -> filename:join([PrivDir, Path])
|
||||
end
|
||||
end,
|
||||
error_logger:info_report(io_lib:format("XdbFile:~s~n", [XdbFileName])),
|
||||
{ok, IoDevice} = file:open(XdbFileName, [read, binary]),
|
||||
load_vector_index(IoDevice),
|
||||
{ok, #state{xdb_fd = IoDevice}}.
|
||||
{ok, HeaderBin} = file:read(IoDevice, ?XDB_HEADER_SIZE),
|
||||
{ok, Header} = ip2region_xdb:parse_header(HeaderBin),
|
||||
Version = resolve_version(Header),
|
||||
SegmentIndexSize = ip2region_xdb:segment_index_size(Version),
|
||||
load_vector_index(IoDevice, Version),
|
||||
{ok, #state{
|
||||
xdb_fd = IoDevice,
|
||||
version = Version,
|
||||
segment_index_size = SegmentIndexSize
|
||||
}}.
|
||||
|
||||
handle_call(Request, From, State) ->
|
||||
try
|
||||
|
|
@ -85,22 +105,19 @@ handle_info(Info, State) ->
|
|||
terminate(_Reason, State) ->
|
||||
#state{xdb_fd = XdbFd} = State,
|
||||
case is_pid(XdbFd) of
|
||||
true ->
|
||||
file:close(XdbFd);
|
||||
_ ->
|
||||
skip
|
||||
true -> file:close(XdbFd);
|
||||
_ -> skip
|
||||
end,
|
||||
ok.
|
||||
|
||||
code_change(_OldVsn, State, _Extra) ->
|
||||
{ok, State}.
|
||||
|
||||
|
||||
%%==========================================
|
||||
%% Internal function
|
||||
%% =========================================
|
||||
do_call({search, Ip}, _From, #state{xdb_fd = IoDevice} = State) ->
|
||||
Reply = search_ip(IoDevice, Ip),
|
||||
Reply = search_ip(IoDevice, Ip, State),
|
||||
{reply, Reply, State};
|
||||
|
||||
do_call(stop, _From, State) ->
|
||||
|
|
@ -118,66 +135,117 @@ do_info(Info, State) ->
|
|||
error_logger:error_report(io:format("unknown info: ~p", [Info])),
|
||||
{noreply, State}.
|
||||
|
||||
|
||||
load_vector_index(IoDevice) ->
|
||||
Key = ip2region_header_loaded,
|
||||
case persistent_term:get(Key, false) of
|
||||
true -> ok;
|
||||
_ ->
|
||||
{ok, <<_Header:?XDB_HEADER_SIZE/binary, VectorIndexBin/binary>> } =
|
||||
file:read(IoDevice, ?XDB_HEADER_SIZE + ?XDB_VECTOR_INDEX_COUNT*8),
|
||||
load_vector_index_aux(VectorIndexBin, 0),
|
||||
persistent_term:put(Key, true)
|
||||
resolve_version(Header) ->
|
||||
case ip2region_xdb:header_version(Header) of
|
||||
2 -> ipv4;
|
||||
3 ->
|
||||
case ip2region_xdb:header_ip_version(Header) of
|
||||
?IP_VERSION_4 -> ipv4;
|
||||
?IP_VERSION_6 -> ipv6;
|
||||
_ -> ipv4
|
||||
end;
|
||||
_ -> ipv4
|
||||
end.
|
||||
|
||||
load_vector_index_aux(<<>>, _Index) -> ok;
|
||||
load_vector_index_aux(<<SPtr:32/little, EPtr:32/little, VectorIndexBin/binary>>, Index) ->
|
||||
Term = {Index, SPtr, EPtr},
|
||||
ets:insert(?XDB_VECTOR_INDEX, Term),
|
||||
load_vector_index_aux(VectorIndexBin, Index + 1).
|
||||
load_vector_index(IoDevice, Version) ->
|
||||
Table = vector_index_table(Version),
|
||||
case ets:info(Table, size) of
|
||||
undefined ->
|
||||
Opts = [named_table, set, public, {read_concurrency, true}, {keypos, 1}],
|
||||
ets:new(Table, Opts),
|
||||
load_vector_index_data(IoDevice, Table);
|
||||
0 ->
|
||||
load_vector_index_data(IoDevice, Table);
|
||||
_ ->
|
||||
ok
|
||||
end.
|
||||
|
||||
load_vector_index_data(IoDevice, Table) ->
|
||||
{ok, VectorIndexBin} =
|
||||
file:read(IoDevice, ?XDB_VECTOR_INDEX_COUNT * 8),
|
||||
load_vector_index_aux(VectorIndexBin, 0, Table).
|
||||
|
||||
search_ip(IoDevice, Ip) ->
|
||||
IntIp = ip2region_util:ipv4_to_n(Ip),
|
||||
case ets:lookup(?IP2REGION_CACHE, IntIp) of
|
||||
[{_IntIp, RegionInfo}] ->
|
||||
load_vector_index_aux(<<>>, _Index, _Table) -> ok;
|
||||
load_vector_index_aux(<<SPtr:32/little, EPtr:32/little, VectorIndexBin/binary>>, Index, Table) ->
|
||||
ets:insert(Table, {Index, SPtr, EPtr}),
|
||||
load_vector_index_aux(VectorIndexBin, Index + 1, Table).
|
||||
|
||||
search_ip(IoDevice, IpInt, State) when is_integer(IpInt) ->
|
||||
search_ip(IoDevice, <<IpInt:32>>, State);
|
||||
search_ip(IoDevice, Ip, #state{version = Version, segment_index_size = SegSize}) ->
|
||||
CacheTable = cache_table(Version),
|
||||
VectorTable = vector_index_table(Version),
|
||||
SegmentTable = segment_index_table(Version),
|
||||
case ets:lookup(CacheTable, Ip) of
|
||||
[{_, RegionInfo}] ->
|
||||
RegionInfo;
|
||||
_ ->
|
||||
<<A:8, B:8, _Rest/binary>> = <<IntIp:32>>,
|
||||
<<A:8, B:8, _/binary>> = Ip,
|
||||
VectorIdx = A * ?XDB_VECTOR_COLS + B,
|
||||
[{_, SPtr, EPtr}] = ets:lookup(?XDB_VECTOR_INDEX, VectorIdx),
|
||||
RegionInfo = search_ip(IoDevice, IntIp, SPtr, EPtr, 0, (EPtr - SPtr) div ?XDB_SEGMENT_INDEX_SIZE),
|
||||
ets:insert_new(?IP2REGION_CACHE, {IntIp, RegionInfo}),
|
||||
[{_, SPtr, EPtr}] = ets:lookup(VectorTable, VectorIdx),
|
||||
RegionInfo = search_ip(IoDevice, Ip, SPtr, EPtr, 0,
|
||||
(EPtr - SPtr) div SegSize, SegSize, Version, SegmentTable),
|
||||
ets:insert_new(CacheTable, {Ip, RegionInfo}),
|
||||
RegionInfo
|
||||
end.
|
||||
|
||||
search_ip(IoDevice, IntIp, SPtr, EPtr, Low, High) when Low =< High ->
|
||||
search_ip(IoDevice, Ip, SPtr, EPtr, Low, High, SegSize, Version, SegmentTable) when Low =< High ->
|
||||
Middle = (Low + High) bsr 1,
|
||||
SPtr2 = SPtr + Middle * ?XDB_SEGMENT_INDEX_SIZE,
|
||||
{SIp, EIp, DataLen, DataPtr} = read_segement_index(IoDevice, SPtr2),
|
||||
if
|
||||
IntIp < SIp ->
|
||||
search_ip(IoDevice, IntIp, SPtr, EPtr, Low, Middle - 1);
|
||||
IntIp > EIp ->
|
||||
search_ip(IoDevice, IntIp, SPtr, EPtr, Middle + 1, High);
|
||||
true ->
|
||||
SPtr2 = SPtr + Middle * SegSize,
|
||||
{SIp, EIp, DataLen, DataPtr} = read_segment_index(IoDevice, SPtr2, SegSize, SegmentTable),
|
||||
case ip_in_range(Ip, SIp, EIp, Version) of
|
||||
below ->
|
||||
search_ip(IoDevice, Ip, SPtr, EPtr, Low, Middle - 1, SegSize, Version, SegmentTable);
|
||||
above ->
|
||||
search_ip(IoDevice, Ip, SPtr, EPtr, Middle + 1, High, SegSize, Version, SegmentTable);
|
||||
inside ->
|
||||
{ok, DataBin} = read_file(IoDevice, DataPtr, DataLen),
|
||||
unicode:characters_to_nfc_list(DataBin)
|
||||
end;
|
||||
search_ip(_IoDevice, _IntIp, _SPtr, _EPtr, _Low, _High) ->
|
||||
search_ip(_IoDevice, _Ip, _SPtr, _EPtr, _Low, _High, _SegSize, _Version, _SegmentTable) ->
|
||||
{error, unknown}.
|
||||
|
||||
ip_in_range(Ip, SIp, EIp, ipv4) ->
|
||||
<<InputInt:32>> = Ip,
|
||||
if
|
||||
InputInt < SIp -> below;
|
||||
InputInt > EIp -> above;
|
||||
true -> inside
|
||||
end;
|
||||
ip_in_range(Ip, SIp, EIp, ipv6) ->
|
||||
if
|
||||
Ip < SIp -> below;
|
||||
Ip > EIp -> above;
|
||||
true -> inside
|
||||
end.
|
||||
|
||||
read_file(IoDevice, Position, DataLength) ->
|
||||
file:position(IoDevice, {bof, Position}),
|
||||
file:read(IoDevice, DataLength).
|
||||
|
||||
read_segement_index(IoDevice, SPtr) ->
|
||||
case ets:lookup(?XDB_SEGMENT_INDEX, SPtr) of
|
||||
read_segment_index(IoDevice, SPtr, SegSize, SegmentTable) ->
|
||||
case ets:lookup(SegmentTable, SPtr) of
|
||||
[{_SPtr, SIp, EIp, DataLen, DataPtr}] ->
|
||||
{SIp, EIp, DataLen, DataPtr};
|
||||
_ ->
|
||||
{ok, <<SIp:32/little, EIp:32/little, DataLen:16/little, DataPtr:32/little>>} =
|
||||
read_file(IoDevice, SPtr, ?XDB_SEGMENT_INDEX_SIZE),
|
||||
ets:insert_new(?XDB_SEGMENT_INDEX, {SPtr, SIp, EIp, DataLen, DataPtr}),
|
||||
{ok, Bin} = read_file(IoDevice, SPtr, SegSize),
|
||||
{SIp, EIp, DataLen, DataPtr} = decode_segment_index(Bin, SegSize),
|
||||
ets:insert_new(SegmentTable, {SPtr, SIp, EIp, DataLen, DataPtr}),
|
||||
{SIp, EIp, DataLen, DataPtr}
|
||||
end.
|
||||
|
||||
decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V4) ->
|
||||
<<SIp:32/little, EIp:32/little, DataLen:16/little, DataPtr:32/little>> = Bin,
|
||||
{SIp, EIp, DataLen, DataPtr};
|
||||
decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V6) ->
|
||||
<<SIp:16/binary, EIp:16/binary, DataLen:16/little, DataPtr:32/little>> = Bin,
|
||||
{SIp, EIp, DataLen, DataPtr}.
|
||||
|
||||
vector_index_table(ipv4) -> ?XDB_VECTOR_INDEX_V4;
|
||||
vector_index_table(ipv6) -> ?XDB_VECTOR_INDEX_V6.
|
||||
|
||||
segment_index_table(ipv4) -> ?XDB_SEGMENT_INDEX_V4;
|
||||
segment_index_table(ipv6) -> ?XDB_SEGMENT_INDEX_V6.
|
||||
|
||||
cache_table(ipv4) -> ?IP2REGION_CACHE_V4;
|
||||
cache_table(ipv6) -> ?IP2REGION_CACHE_V6.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,13 @@
|
|||
-module(ip2region_worker_test).
|
||||
-include_lib("eunit/include/eunit.hrl").
|
||||
-include("ip2region.hrl").
|
||||
|
||||
worker_search_v4_binary_test() ->
|
||||
ip2region_sup:create_table(),
|
||||
{ok, Pid} = ip2region_worker:start_link([{xdb_file, "ip2region.xdb"}]),
|
||||
try
|
||||
Region = ip2region_worker:search(Pid, <<1,0,8,0>>),
|
||||
?assert(is_list(Region))
|
||||
after
|
||||
ip2region_worker:stop(Pid)
|
||||
end.
|
||||
Loading…
Reference in New Issue