This commit is contained in:
anlicheng 2025-12-17 00:01:43 +08:00
parent 46a4d1e015
commit 9da280de08
2 changed files with 31 additions and 47 deletions

View File

@ -1,6 +1,6 @@
-module(dns_pending_wheel). -module(dns_pending_wheel).
-export([start/0, put/2, get/1, delete/1]). -export([start/0, insert/2, lookup/1, delete/1]).
-define(TTL, 5). -define(TTL, 5).
-define(TICK_MS, 1000). -define(TICK_MS, 1000).
@ -11,26 +11,21 @@
%%% ===================================================== %%% =====================================================
start() -> start() ->
ets:new(dns_pending_data, [ordered_set, public, named_table]), ets:new(dns_pending_data, [ordered_set, public, named_table, {read_concurrency, true}, {write_concurrency, true}]),
ets:new(dns_pending_wheel, [bag, public, named_table]), ets:new(dns_pending_wheel, [bag, public, named_table, {read_concurrency, true}, {write_concurrency, true}]),
start_scanner(). start_scanner().
-spec put(Key :: any(), Val :: any()) -> ok. -spec insert(Key :: any(), Val :: any()) -> ok.
put(Key, Val) -> insert(Key, Val) ->
Tick = now_tick(), Tick = now_tick(),
Slot = Tick rem ?WHEEL_SIZE, Slot = Tick rem ?WHEEL_SIZE,
ets:insert(dns_pending_data, {Key, {Val, Tick}}), ets:insert(dns_pending_data, {Key, {Val, Tick}}),
ets:insert(dns_pending_wheel, {Slot, {Key, Tick}}), ets:insert(dns_pending_wheel, {Slot, {Key, Tick}}),
ok. ok.
-spec get(Key :: any()) -> error | {ok, Val :: any()}. -spec lookup(Key :: any()) -> [term()].
get(Key) -> lookup(Key) ->
case ets:lookup(dns_pending_data, Key) of ets:lookup(dns_pending_data, Key).
[{Key, {Val, _Tick}}] ->
{ok, Val};
[] ->
error
end.
-spec delete(Key :: any()) -> ok. -spec delete(Key :: any()) -> ok.
delete(Key) -> delete(Key) ->

View File

@ -21,11 +21,10 @@
-export([forward/5]). -export([forward/5]).
-define(SERVER, ?MODULE). -define(SERVER, ?MODULE).
-define(REQUEST_TTL, 5000).
-record(state, { -record(state, {
socket, socket,
tid, idx :: integer(),
dns_servers = [] dns_servers = []
}). }).
@ -53,12 +52,9 @@ start_link(Args) when is_list(Args) ->
{stop, Reason :: term()} | ignore). {stop, Reason :: term()} | ignore).
init([]) -> init([]) ->
{ok, DnsServers} = application:get_env(sdlan, public_dns_servers), {ok, DnsServers} = application:get_env(sdlan, public_dns_servers),
{ok, Sock} = gen_udp:open(0, [binary, {active, true}]), {ok, Sock} = gen_udp:open(0, [binary, {active, true}]),
%% ets来保存映射关系 Idx = erlang:unique_integer([monotonic, positive]),
Tid = ets:new(random_table(), [set, {read_concurrency, true}, {write_concurrency, true}, private]), {ok, #state{socket = Sock, idx = Idx, dns_servers = DnsServers}}.
{ok, #state{socket = Sock, tid = Tid, dns_servers = DnsServers}}.
%% @private %% @private
%% @doc Handling call messages %% @doc Handling call messages
@ -79,15 +75,15 @@ handle_call(_Request, _From, State = #state{}) ->
{noreply, NewState :: #state{}} | {noreply, NewState :: #state{}} |
{noreply, NewState :: #state{}, timeout() | hibernate} | {noreply, NewState :: #state{}, timeout() | hibernate} |
{stop, Reason :: term(), NewState :: #state{}}). {stop, Reason :: term(), NewState :: #state{}}).
handle_cast({forward, ReceiverPid, Ref, Request, #dns_message{id = TxId, questions = [#dns_query{name = QName, type = QType, class = QClass}|_]}}, State = #state{socket = Socket, tid = Tid, dns_servers = DnsServers}) -> handle_cast({forward, ReceiverPid, Ref, Request, #dns_message{id = TxId, questions = [#dns_query{name = QName, type = QType, class = QClass}|_]}},
Keys = lists:foldl(fun({DnsIp, DnsPort}, Acc) -> State = #state{socket = Socket, idx = Idx, dns_servers = DnsServers}) ->
lists:foreach(fun({DnsIp, DnsPort}) ->
ok = gen_udp:send(Socket, DnsIp, DnsPort, Request), ok = gen_udp:send(Socket, DnsIp, DnsPort, Request),
Key = {TxId, DnsIp, DnsPort, QName, QType, QClass}, Key = {Idx, TxId, DnsIp, DnsPort, QName, QType, QClass},
lager:debug("[dns_resolver] key: ~p, send to: ~p, packet: ~p", [Key, {DnsIp, DnsPort}, Request]), lager:debug("[dns_resolver] key: ~p, send to: ~p, packet: ~p", [Key, {DnsIp, DnsPort}, Request]),
true = ets:insert(Tid, {Key, Ref, ReceiverPid}), dns_pending_wheel:insert(Key, {Ref, ReceiverPid})
[Key|Acc] end, DnsServers),
end, [], DnsServers),
erlang:start_timer(?REQUEST_TTL, self(), {clean_ticker, Keys}),
{noreply, State}. {noreply, State}.
@ -97,21 +93,18 @@ handle_cast({forward, ReceiverPid, Ref, Request, #dns_message{id = TxId, questio
{noreply, NewState :: #state{}} | {noreply, NewState :: #state{}} |
{noreply, NewState :: #state{}, timeout() | hibernate} | {noreply, NewState :: #state{}, timeout() | hibernate} |
{stop, Reason :: term(), NewState :: #state{}}). {stop, Reason :: term(), NewState :: #state{}}).
handle_info({udp, Socket, TargetIp, TargetPort, Resp}, State = #state{tid = Tid, socket = Socket}) -> handle_info({udp, Socket, TargetIp, TargetPort, Resp}, State = #state{socket = Socket, idx = Idx}) ->
try dns:decode_message(Resp) of try dns:decode_message(Resp) of
#dns_message{id = TxId, questions = [#dns_query{name = QName, type = QType, class = QClass}|_]} -> #dns_message{id = TxId, questions = [#dns_query{name = QName, type = QType, class = QClass}|_]} ->
Key = {TxId, TargetIp, TargetPort, QName, QType, QClass}, Key = {Idx, TxId, TargetIp, TargetPort, QName, QType, QClass},
Records = ets:take(Tid, Key), Records = dns_pending_wheel:lookup(Key),
dns_pending_wheel:delete(Key),
resolver_reply(Records, Resp); resolver_reply(Records, Resp);
_ -> _ ->
ok ok
catch error:_ -> catch error:_ ->
ok ok
end, end,
{noreply, State};
handle_info({timeout, _, {clean_ticker, Keys}}, State = #state{tid = Tid}) ->
lists:foreach(fun(Key) -> ets:delete(Tid, Key) end, Keys),
{noreply, State}. {noreply, State}.
%% @private %% @private
@ -136,17 +129,13 @@ code_change(_OldVsn, State = #state{}, _Extra) ->
%%% Internal functions %%% Internal functions
%%%=================================================================== %%%===================================================================
-spec random_table() -> atom().
random_table() ->
list_to_atom("udp_ets:" ++ integer_to_list(erlang:unique_integer([monotonic, positive]))).
-spec resolver_reply(list(), Resp :: binary()) -> no_return(). -spec resolver_reply(list(), Resp :: binary()) -> no_return().
resolver_reply([{_, Ref, ReceiverPid}], Resp) when is_binary(Resp) -> resolver_reply(Records, Resp) when is_binary(Resp) ->
case is_process_alive(ReceiverPid) of lists:foreach(fun({_, {Ref, ReceiverPid}}) ->
true -> case is_process_alive(ReceiverPid) of
ReceiverPid ! {dns_resolver_reply, Ref, Resp}; true ->
false -> ReceiverPid ! {dns_resolver_reply, Ref, Resp};
ok false ->
end; ok
resolver_reply(_, _) -> end
ok. end, Records).