Add the new 'emqx_limiter' module

This commit is contained in:
Feng Lee 2019-11-01 08:10:58 +08:00
parent 6f30dca4ba
commit 2b1b58fc66
8 changed files with 611 additions and 285 deletions

View File

@ -1,4 +1,4 @@
%%-------------------------------------------------------------------- %%-------------------------------------------------------------------
%% Copyright (c) 2019 EMQ Technologies Co., Ltd. All Rights Reserved. %% Copyright (c) 2019 EMQ Technologies Co., Ltd. All Rights Reserved.
%% %%
%% Licensed under the Apache License, Version 2.0 (the "License"); %% Licensed under the Apache License, Version 2.0 (the "License");
@ -93,7 +93,6 @@ start_link() ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% @doc Register a channel. %% @doc Register a channel.
%% Channel will be unregistered automatically when the channel process dies
-spec(register_channel(emqx_types:clientid()) -> ok). -spec(register_channel(emqx_types:clientid()) -> ok).
register_channel(ClientId) when is_binary(ClientId) -> register_channel(ClientId) when is_binary(ClientId) ->
register_channel(ClientId, self()). register_channel(ClientId, self()).

View File

@ -14,7 +14,7 @@
%% limitations under the License. %% limitations under the License.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% MQTT/TCP Connection %% MQTT/TCP|TLS Connection
-module(emqx_connection). -module(emqx_connection).
-include("emqx.hrl"). -include("emqx.hrl").
@ -40,7 +40,7 @@
-export([call/2]). -export([call/2]).
%% callback %% Callback
-export([init/4]). -export([init/4]).
%% Sys callbacks %% Sys callbacks
@ -50,12 +50,15 @@
, system_get_state/1 , system_get_state/1
]). ]).
%% Internal callbacks %% Internal callback
-export([wakeup_from_hib/2]). -export([wakeup_from_hib/3]).
-import(emqx_misc,
[ maybe_apply/2
, start_timer/2
]).
-record(state, { -record(state, {
%% Parent
parent :: pid(),
%% TCP/TLS Transport %% TCP/TLS Transport
transport :: esockd:transport(), transport :: esockd:transport(),
%% TCP/TLS Socket %% TCP/TLS Socket
@ -64,34 +67,37 @@
peername :: emqx_types:peername(), peername :: emqx_types:peername(),
%% Sockname of the connection %% Sockname of the connection
sockname :: emqx_types:peername(), sockname :: emqx_types:peername(),
%% Sock state %% Sock State
sockstate :: emqx_types:sockstate(), sockstate :: emqx_types:sockstate(),
%% The {active, N} option %% The {active, N} option
active_n :: pos_integer(), active_n :: pos_integer(),
%% Publish Limit %% Limiter
pub_limit :: maybe(esockd_rate_limit:bucket()), limiter :: maybe(emqx_limiter:limiter()),
%% Rate Limit
rate_limit :: maybe(esockd_rate_limit:bucket()),
%% Limit Timer %% Limit Timer
limit_timer :: maybe(reference()), limit_timer :: maybe(reference()),
%% Parser State %% Parse State
parse_state :: emqx_frame:parse_state(), parse_state :: emqx_frame:parse_state(),
%% Serialize function %% Serialize function
serialize :: emqx_frame:serialize_fun(), serialize :: emqx_frame:serialize_fun(),
%% Channel State %% Channel State
channel :: emqx_channel:channel(), channel :: emqx_channel:channel(),
%% Idle timer %% GC State
idle_timer :: reference() gc_state :: maybe(emqx_gc:gc_state()),
%% Stats Timer
stats_timer :: disabled | maybe(reference()),
%% Idle Timer
idle_timer :: maybe(reference())
}). }).
-type(state() :: #state{}). -type(state() :: #state{}).
-define(ACTIVE_N, 100). -define(ACTIVE_N, 100).
-define(INFO_KEYS, [socktype, peername, sockname, sockstate, active_n, -define(INFO_KEYS, [socktype, peername, sockname, sockstate, active_n, limiter]).
pub_limit, rate_limit]).
-define(CONN_STATS, [recv_pkt, recv_msg, send_pkt, send_msg]). -define(CONN_STATS, [recv_pkt, recv_msg, send_pkt, send_msg]).
-define(SOCK_STATS, [recv_oct, recv_cnt, send_oct, send_cnt, send_pend]). -define(SOCK_STATS, [recv_oct, recv_cnt, send_oct, send_cnt, send_pend]).
-define(ENABLED(X), (X =/= undefined)).
-spec(start_link(esockd:transport(), esockd:socket(), proplists:proplist()) -spec(start_link(esockd:transport(), esockd:socket(), proplists:proplist())
-> {ok, pid()}). -> {ok, pid()}).
start_link(Transport, Socket, Options) -> start_link(Transport, Socket, Options) ->
@ -123,13 +129,8 @@ info(sockstate, #state{sockstate = SockSt}) ->
SockSt; SockSt;
info(active_n, #state{active_n = ActiveN}) -> info(active_n, #state{active_n = ActiveN}) ->
ActiveN; ActiveN;
info(pub_limit, #state{pub_limit = PubLimit}) -> info(limiter, #state{limiter = Limiter}) ->
limit_info(PubLimit); maybe_apply(fun emqx_limiter:info/1, Limiter).
info(rate_limit, #state{rate_limit = RateLimit}) ->
limit_info(RateLimit).
limit_info(Limit) ->
emqx_misc:maybe_apply(fun esockd_rate_limit:info/1, Limit).
%% @doc Get stats of the connection/channel. %% @doc Get stats of the connection/channel.
-spec(stats(pid()|state()) -> emqx_types:stats()). -spec(stats(pid()|state()) -> emqx_types:stats()).
@ -147,6 +148,13 @@ stats(#state{transport = Transport,
ProcStats = emqx_misc:proc_stats(), ProcStats = emqx_misc:proc_stats(),
lists:append([SockStats, ConnStats, ChanStats, ProcStats]). lists:append([SockStats, ConnStats, ChanStats, ProcStats]).
attrs(#state{active_n = ActiveN, sockstate = SockSt, channel = Channel}) ->
SockAttrs = #{active_n => ActiveN,
sockstate => SockSt
},
ChanAttrs = emqx_channel:attrs(Channel),
maps:merge(ChanAttrs, #{sockinfo => SockAttrs}).
call(Pid, Req) -> call(Pid, Req) ->
gen_server:call(Pid, Req, infinity). gen_server:call(Pid, Req, infinity).
@ -169,7 +177,6 @@ init(Parent, Transport, RawSocket, Options) ->
do_init(Parent, Transport, Socket, Options) -> do_init(Parent, Transport, Socket, Options) ->
{ok, Peername} = Transport:ensure_ok_or_exit(peername, [Socket]), {ok, Peername} = Transport:ensure_ok_or_exit(peername, [Socket]),
{ok, Sockname} = Transport:ensure_ok_or_exit(sockname, [Socket]), {ok, Sockname} = Transport:ensure_ok_or_exit(sockname, [Socket]),
emqx_logger:set_metadata_peername(esockd_net:format(Peername)),
Peercert = Transport:ensure_ok_or_exit(peercert, [Socket]), Peercert = Transport:ensure_ok_or_exit(peercert, [Socket]),
ConnInfo = #{socktype => Transport:type(Socket), ConnInfo = #{socktype => Transport:type(Socket),
peername => Peername, peername => Peername,
@ -179,42 +186,39 @@ do_init(Parent, Transport, Socket, Options) ->
}, },
Zone = proplists:get_value(zone, Options), Zone = proplists:get_value(zone, Options),
ActiveN = proplists:get_value(active_n, Options, ?ACTIVE_N), ActiveN = proplists:get_value(active_n, Options, ?ACTIVE_N),
PubLimit = init_limiter(emqx_zone:get_env(Zone, publish_limit)), Limiter = emqx_limiter:init(Options),
RateLimit = init_limiter(proplists:get_value(rate_limit, Options)), FrameOpts = emqx_zone:mqtt_frame_options(Zone),
FrameOpts = emqx_zone:frame_options(Zone),
ParseState = emqx_frame:initial_parse_state(FrameOpts), ParseState = emqx_frame:initial_parse_state(FrameOpts),
Serialize = emqx_frame:serialize_fun(), Serialize = emqx_frame:serialize_fun(),
Channel = emqx_channel:init(ConnInfo, Options), Channel = emqx_channel:init(ConnInfo, Options),
IdleTimout = emqx_zone:get_env(Zone, idle_timeout, 30000), GcState = emqx_zone:init_gc_state(Zone),
IdleTimer = emqx_misc:start_timer(IdleTimout, idle_timeout), StatsTimer = emqx_zone:stats_timer(Zone),
HibAfterTimeout = emqx_zone:get_env(Zone, hibernate_after, IdleTimout*2), IdleTimeout = emqx_zone:idle_timeout(Zone),
State = #state{parent = Parent, IdleTimer = start_timer(IdleTimeout, idle_timeout),
transport = Transport, emqx_misc:tune_heap_size(emqx_zone:oom_policy(Zone)),
emqx_logger:set_metadata_peername(esockd_net:format(Peername)),
State = #state{transport = Transport,
socket = Socket, socket = Socket,
peername = Peername, peername = Peername,
sockname = Sockname, sockname = Sockname,
sockstate = idle, sockstate = idle,
active_n = ActiveN, active_n = ActiveN,
pub_limit = PubLimit, limiter = Limiter,
rate_limit = RateLimit,
parse_state = ParseState, parse_state = ParseState,
serialize = Serialize, serialize = Serialize,
channel = Channel, channel = Channel,
gc_state = GcState,
stats_timer = StatsTimer,
idle_timer = IdleTimer idle_timer = IdleTimer
}, },
case activate_socket(State) of case activate_socket(State) of
{ok, NState} -> {ok, NState} ->
hibernate(NState, #{hibernate_after => HibAfterTimeout}); hibernate(Parent, NState, #{idle_timeout => IdleTimeout});
{error, Reason} -> {error, Reason} ->
ok = Transport:fast_close(Socket), ok = Transport:fast_close(Socket),
exit_on_sock_error(Reason) exit_on_sock_error(Reason)
end. end.
-compile({inline, [init_limiter/1]}).
init_limiter(undefined) -> undefined;
init_limiter({Rate, Burst}) ->
esockd_rate_limit:new(Rate, Burst).
exit_on_sock_error(Reason) when Reason =:= einval; exit_on_sock_error(Reason) when Reason =:= einval;
Reason =:= enotconn; Reason =:= enotconn;
Reason =:= closed -> Reason =:= closed ->
@ -227,8 +231,7 @@ exit_on_sock_error(Reason) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Recv Loop %% Recv Loop
recvloop(State = #state{parent = Parent}, recvloop(Parent, State, Options = #{idle_timeout := IdleTimeout}) ->
Options = #{hibernate_after := HibAfterTimeout}) ->
receive receive
{system, From, Request} -> {system, From, Request} ->
sys:handle_system_msg(Request, From, Parent, sys:handle_system_msg(Request, From, Parent,
@ -236,33 +239,49 @@ recvloop(State = #state{parent = Parent},
{'EXIT', Parent, Reason} -> {'EXIT', Parent, Reason} ->
terminate(Reason, State); terminate(Reason, State);
Msg -> Msg ->
process_msg([Msg], State, Options) NState = ensure_stats_timer(IdleTimeout, State),
process_msg([Msg], Parent, NState, Options)
after after
HibAfterTimeout -> IdleTimeout ->
hibernate(State, Options) NState = cancel_stats_timer(State),
hibernate(Parent, NState, Options)
end. end.
hibernate(State, Options) -> hibernate(Parent, State, Options) ->
proc_lib:hibernate(?MODULE, wakeup_from_hib, [State, Options]). proc_lib:hibernate(?MODULE, wakeup_from_hib, [Parent, State, Options]).
wakeup_from_hib(State, Options) -> wakeup_from_hib(Parent, State, Options) ->
%% Maybe do something later here. %% Maybe do something later here.
recvloop(State, Options). recvloop(Parent, State, Options).
%%--------------------------------------------------------------------
%% Ensure/cancel stats timer
-compile({inline, [ensure_stats_timer/2]}).
ensure_stats_timer(Timeout, State = #state{stats_timer = undefined}) ->
State#state{stats_timer = start_timer(Timeout, emit_stats)};
ensure_stats_timer(_Timeout, State) -> State.
-compile({inline, [cancel_stats_timer/1]}).
cancel_stats_timer(State = #state{stats_timer = TRef}) when is_reference(TRef) ->
ok = emqx_misc:cancel_timer(TRef),
State#state{stats_timer = undefined};
cancel_stats_timer(State) -> State.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Process next Msg %% Process next Msg
process_msg([], State, Options) -> process_msg([], Parent, State, Options) ->
recvloop(State, Options); recvloop(Parent, State, Options);
process_msg([Msg|More], State, Options) -> process_msg([Msg|More], Parent, State, Options) ->
case catch handle_msg(Msg, State) of case catch handle_msg(Msg, State) of
ok -> ok ->
process_msg(More, State, Options); process_msg(More, Parent, State, Options);
{ok, NState} -> {ok, NState} ->
process_msg(More, NState, Options); process_msg(More, Parent, NState, Options);
{ok, NextMsgs, NState} -> {ok, Msgs, NState} ->
process_msg(append_msg(NextMsgs, More), NState, Options); process_msg(append_msg(Msgs, More), Parent, NState, Options);
{stop, Reason} -> {stop, Reason} ->
terminate(Reason, State); terminate(Reason, State);
{stop, Reason, NState} -> {stop, Reason, NState} ->
@ -284,14 +303,12 @@ handle_msg({'$gen_call', From, Req}, State) ->
stop(Reason, NState) stop(Reason, NState)
end; end;
handle_msg({Inet, _Sock, Data}, State = #state{channel = Channel}) handle_msg({Inet, _Sock, Data}, State) when Inet == tcp; Inet == ssl ->
when Inet == tcp; Inet == ssl ->
?LOG(debug, "RECV ~p", [Data]), ?LOG(debug, "RECV ~p", [Data]),
Oct = iolist_size(Data), Oct = iolist_size(Data),
emqx_pd:update_counter(incoming_bytes, Oct), emqx_pd:inc_counter(incoming_bytes, Oct),
ok = emqx_metrics:inc('bytes.received', Oct), ok = emqx_metrics:inc('bytes.received', Oct),
NChannel = emqx_channel:recvd(Oct, Channel), parse_incoming(Data, State);
parse_incoming(Data, State#state{channel = NChannel});
handle_msg({incoming, Packet = ?CONNECT_PACKET(ConnPkt)}, handle_msg({incoming, Packet = ?CONNECT_PACKET(ConnPkt)},
State = #state{idle_timer = IdleTimer}) -> State = #state{idle_timer = IdleTimer}) ->
@ -302,6 +319,9 @@ handle_msg({incoming, Packet = ?CONNECT_PACKET(ConnPkt)},
}, },
handle_incoming(Packet, NState); handle_incoming(Packet, NState);
handle_msg({incoming, ?PACKET(?PINGREQ)}, State) ->
handle_outgoing(?PACKET(?PINGRESP), State);
handle_msg({incoming, Packet}, State) -> handle_msg({incoming, Packet}, State) ->
handle_incoming(Packet, State); handle_incoming(Packet, State);
@ -315,30 +335,34 @@ handle_msg({Closed, _Sock}, State)
handle_msg({Passive, _Sock}, State) handle_msg({Passive, _Sock}, State)
when Passive == tcp_passive; Passive == ssl_passive -> when Passive == tcp_passive; Passive == ssl_passive ->
%% Rate limit here:) InStats = #{cnt => emqx_pd:reset_counter(incoming_pubs),
NState = ensure_rate_limit(State), oct => emqx_pd:reset_counter(incoming_bytes)
handle_info(activate_socket, NState); },
%% Ensure Rate Limit
%% Rate limit timer expired. NState = ensure_rate_limit(InStats, State),
handle_msg(activate_socket, State) -> %% Run GC and Check OOM
NState = State#state{sockstate = idle, NState1 = check_oom(run_gc(InStats, NState)),
limit_timer = undefined handle_info(activate_socket, NState1);
},
handle_info(activate_socket, NState);
handle_msg(Deliver = {deliver, _Topic, _Msg}, handle_msg(Deliver = {deliver, _Topic, _Msg},
State = #state{channel = Channel}) -> State = #state{channel = Channel}) ->
Delivers = emqx_misc:drain_deliver([Deliver]), Delivers = [Deliver|emqx_misc:drain_deliver()],
Result = emqx_channel:handle_out(Delivers, Channel), Ret = emqx_channel:handle_out(Delivers, Channel),
handle_return(Result, State); handle_chan_return(Ret, State);
handle_msg({outgoing, Packets}, State) -> handle_msg({outgoing, Packets}, State) ->
NState = handle_outgoing(Packets, State), handle_outgoing(Packets, State);
{ok, NState};
%% something sent %% Something sent
handle_msg({inet_reply, _Sock, ok}, _State) -> handle_msg({inet_reply, _Sock, ok}, State = #state{active_n = ActiveN}) ->
ok; case emqx_pd:get_counter(outgoing_pubs) > ActiveN of
true ->
OutStats = #{cnt => emqx_pd:reset_counter(outgoing_pubs),
oct => emqx_pd:reset_counter(outgoing_bytes)
},
{ok, check_oom(run_gc(OutStats, State))};
false -> ok
end;
handle_msg({inet_reply, _Sock, {error, Reason}}, State) -> handle_msg({inet_reply, _Sock, {error, Reason}}, State) ->
handle_info({sock_error, Reason}, State); handle_info({sock_error, Reason}, State);
@ -349,7 +373,8 @@ handle_msg({timeout, TRef, TMsg}, State) ->
handle_msg(Shutdown = {shutdown, _Reason}, State) -> handle_msg(Shutdown = {shutdown, _Reason}, State) ->
stop(Shutdown, State); stop(Shutdown, State);
handle_msg(Msg, State) -> handle_info(Msg, State). handle_msg(Msg, State) ->
handle_info(Msg, State).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Terminate %% Terminate
@ -363,8 +388,8 @@ terminate(Reason, State = #state{channel = Channel}) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Sys callbacks %% Sys callbacks
system_continue(_Parent, _Deb, {State, Options}) -> system_continue(Parent, _Deb, {State, Options}) ->
recvloop(State, Options). recvloop(Parent, State, Options).
system_terminate(Reason, _Parent, _Deb, {State, _}) -> system_terminate(Reason, _Parent, _Deb, {State, _}) ->
terminate(Reason, State). terminate(Reason, State).
@ -392,8 +417,8 @@ handle_call(_From, Req, State = #state{channel = Channel}) ->
shutdown(Reason, Reply, State#state{channel = NChannel}); shutdown(Reason, Reply, State#state{channel = NChannel});
{shutdown, Reason, Reply, OutPacket, NChannel} -> {shutdown, Reason, Reply, OutPacket, NChannel} ->
NState = State#state{channel = NChannel}, NState = State#state{channel = NChannel},
NState1 = handle_outgoing(OutPacket, NState), ok = handle_outgoing(OutPacket, NState),
shutdown(Reason, Reply, NState1) shutdown(Reason, Reply, NState)
end. end.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
@ -402,8 +427,18 @@ handle_call(_From, Req, State = #state{channel = Channel}) ->
handle_timeout(TRef, idle_timeout, State = #state{idle_timer = TRef}) -> handle_timeout(TRef, idle_timeout, State = #state{idle_timer = TRef}) ->
shutdown(idle_timeout, State); shutdown(idle_timeout, State);
handle_timeout(TRef, emit_stats, State) -> handle_timeout(TRef, limit_timeout, State = #state{limit_timer = TRef}) ->
handle_timeout(TRef, {emit_stats, stats(State)}, State); NState = State#state{sockstate = idle,
limit_timer = undefined
},
handle_info(activate_socket, NState);
handle_timeout(TRef, emit_stats, State = #state{stats_timer = TRef,
channel = Channel}) ->
#{clientid := ClientId} = emqx_channel:info(clientinfo, Channel),
(ClientId =/= undefined) andalso
emqx_cm:set_chan_stats(ClientId, stats(State)),
{ok, State#state{stats_timer = undefined}};
handle_timeout(TRef, keepalive, State = #state{transport = Transport, handle_timeout(TRef, keepalive, State = #state{transport = Transport,
socket = Socket}) -> socket = Socket}) ->
@ -415,7 +450,8 @@ handle_timeout(TRef, keepalive, State = #state{transport = Transport,
end; end;
handle_timeout(TRef, Msg, State = #state{channel = Channel}) -> handle_timeout(TRef, Msg, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_timeout(TRef, Msg, Channel), State). Ret = emqx_channel:handle_timeout(TRef, Msg, Channel),
handle_chan_return(Ret, State).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Parse incoming data %% Parse incoming data
@ -450,30 +486,30 @@ next_incoming_msgs(Packets) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle incoming packet %% Handle incoming packet
handle_incoming(Packet = ?PACKET(Type), State = #state{channel = Channel}) -> handle_incoming(Packet, State = #state{channel = Channel})
_ = inc_incoming_stats(Type), when is_record(Packet, mqtt_packet) ->
ok = emqx_metrics:inc_recv(Packet), ok = inc_incoming_stats(Packet),
?LOG(debug, "RECV ~s", [emqx_packet:format(Packet)]), ?LOG(debug, "RECV ~s", [emqx_packet:format(Packet)]),
handle_return(emqx_channel:handle_in(Packet, Channel), State); handle_chan_return(emqx_channel:handle_in(Packet, Channel), State);
handle_incoming(FrameError, State = #state{channel = Channel}) -> handle_incoming(FrameError, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_in(FrameError, Channel), State). handle_chan_return(emqx_channel:handle_in(FrameError, Channel), State).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle channel return %% Handle channel return
handle_return(ok, State) -> handle_chan_return(ok, State) ->
{ok, State}; {ok, State};
handle_return({ok, NChannel}, State) -> handle_chan_return({ok, NChannel}, State) ->
{ok, State#state{channel = NChannel}}; {ok, State#state{channel = NChannel}};
handle_return({ok, Replies, NChannel}, State) -> handle_chan_return({ok, Replies, NChannel}, State) ->
{ok, next_msgs(Replies), State#state{channel = NChannel}}; {ok, next_msgs(Replies), State#state{channel = NChannel}};
handle_return({shutdown, Reason, NChannel}, State) -> handle_chan_return({shutdown, Reason, NChannel}, State) ->
shutdown(Reason, State#state{channel = NChannel}); shutdown(Reason, State#state{channel = NChannel});
handle_return({shutdown, Reason, OutPacket, NChannel}, State) -> handle_chan_return({shutdown, Reason, OutPacket, NChannel}, State) ->
NState = State#state{channel = NChannel}, NState = State#state{channel = NChannel},
NState1 = handle_outgoing(OutPacket, NState), ok = handle_outgoing(OutPacket, NState),
shutdown(Reason, NState1). shutdown(Reason, NState).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle outgoing packets %% Handle outgoing packets
@ -485,14 +521,13 @@ handle_outgoing(Packet, State) ->
send((serialize_and_inc_stats_fun(State))(Packet), State). send((serialize_and_inc_stats_fun(State))(Packet), State).
serialize_and_inc_stats_fun(#state{serialize = Serialize}) -> serialize_and_inc_stats_fun(#state{serialize = Serialize}) ->
fun(Packet = ?PACKET(Type)) -> fun(Packet) ->
case Serialize(Packet) of case Serialize(Packet) of
<<>> -> ?LOG(warning, "~s is discarded due to the frame is too large!", <<>> -> ?LOG(warning, "~s is discarded due to the frame is too large!",
[emqx_packet:format(Packet)]), [emqx_packet:format(Packet)]),
<<>>; <<>>;
Data -> _ = inc_outgoing_stats(Type), Data -> ?LOG(debug, "SEND ~s", [emqx_packet:format(Packet)]),
_ = emqx_metrics:inc_sent(Packet), ok = inc_outgoing_stats(Packet),
?LOG(debug, "SEND ~s", [emqx_packet:format(Packet)]),
Data Data
end end
end. end.
@ -500,52 +535,52 @@ serialize_and_inc_stats_fun(#state{serialize = Serialize}) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Send data %% Send data
send(IoData, State = #state{transport = Transport, -spec(send(iodata(), state()) -> ok).
socket = Socket, send(IoData, #state{transport = Transport, socket = Socket}) ->
channel = Channel}) ->
Oct = iolist_size(IoData), Oct = iolist_size(IoData),
ok = emqx_metrics:inc('bytes.sent', Oct), ok = emqx_metrics:inc('bytes.sent', Oct),
emqx_pd:inc_counter(outgoing_bytes, Oct),
case Transport:async_send(Socket, IoData) of case Transport:async_send(Socket, IoData) of
ok -> ok -> ok;
State#state{channel = emqx_channel:sent(Oct, Channel)};
Error = {error, _Reason} -> Error = {error, _Reason} ->
%% Simulate an inet_reply to postpone handling the error %% Send an inet_reply to postpone handling the error
self() ! {inet_reply, Socket, Error}, State self() ! {inet_reply, Socket, Error},
ok
end. end.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle Info %% Handle Info
handle_info({connack, ConnAck}, State = #state{active_n = ActiveN, handle_info({connack, ConnAck}, State = #state{channel = Channel}) ->
sockstate = SockSt, #{clientid := ClientId} = emqx_channel:info(clientinfo, Channel),
channel = Channel}) -> ok = emqx_cm:register_channel(ClientId),
NState = handle_outgoing(ConnAck, State), ok = emqx_cm:set_chan_attrs(ClientId, attrs(State)),
ChanAttrs = emqx_channel:attrs(Channel), ok = emqx_cm:set_chan_stats(ClientId, stats(State)),
SockAttrs = #{active_n => ActiveN, ok = handle_outgoing(ConnAck, State);
sockstate => SockSt
},
Attrs = maps:merge(ChanAttrs, #{sockinfo => SockAttrs}),
handle_info({register, Attrs, stats(State)}, NState);
handle_info({enter, disconnected}, State = #state{active_n = ActiveN, handle_info({enter, disconnected}, State = #state{channel = Channel}) ->
sockstate = SockSt, #{clientid := ClientId} = emqx_channel:info(clientinfo, Channel),
channel = Channel}) -> emqx_cm:set_chan_attrs(ClientId, attrs(State)),
ChanAttrs = emqx_channel:attrs(Channel), emqx_cm:set_chan_stats(ClientId, stats(State));
SockAttrs = #{active_n => ActiveN,
sockstate => SockSt
},
Attrs = maps:merge(ChanAttrs, #{sockinfo => SockAttrs}),
handle_info({register, Attrs, stats(State)}, State);
handle_info(activate_socket, State) -> handle_info(activate_socket, State = #state{sockstate = OldSst}) ->
case activate_socket(State) of case activate_socket(State) of
{ok, NState} -> {ok, NState}; {ok, NState = #state{sockstate = NewSst}} ->
if OldSst =/= NewSst ->
{ok, {event, sockstate_changed}, NState};
true -> {ok, NState}
end;
{error, Reason} -> {error, Reason} ->
handle_info({sock_error, Reason}, State) handle_info({sock_error, Reason}, State)
end; end;
handle_info({event, sockstate_changed}, State = #state{channel = Channel}) ->
#{clientid := ClientId} = emqx_channel:info(clientinfo, Channel),
ClientId =/= undefined andalso emqx_cm:set_chan_attrs(ClientId, attrs(State));
%%TODO: this is not right %%TODO: this is not right
handle_info({sock_error, _Reason}, #state{sockstate = closed}) -> ok; handle_info({sock_error, _Reason}, #state{sockstate = closed}) ->
ok;
handle_info({sock_error, Reason}, State) -> handle_info({sock_error, Reason}, State) ->
?LOG(debug, "Socket error: ~p", [Reason]), ?LOG(debug, "Socket error: ~p", [Reason]),
handle_info({sock_closed, Reason}, close_socket(State)); handle_info({sock_closed, Reason}, close_socket(State));
@ -560,7 +595,45 @@ handle_info({close, Reason}, State) ->
handle_info({sock_closed, Reason}, close_socket(State)); handle_info({sock_closed, Reason}, close_socket(State));
handle_info(Info, State = #state{channel = Channel}) -> handle_info(Info, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_info(Info, Channel), State). handle_chan_return(emqx_channel:handle_info(Info, Channel), State).
%%--------------------------------------------------------------------
%% Ensure rate limit
ensure_rate_limit(Stats, State = #state{limiter = Limiter}) ->
case ?ENABLED(limiter) andalso emqx_limiter:check(Stats, Limiter) of
false -> State;
{ok, Limiter1} ->
State#state{limiter = Limiter1};
{pause, Time, Limiter1} ->
?LOG(debug, "Pause ~pms due to rate limit", [Time]),
TRef = start_timer(Time, limit_timeout),
State#state{sockstate = blocked,
limiter = Limiter1,
limit_timer = TRef
}
end.
%%--------------------------------------------------------------------
%% Run GC and Check OOM
run_gc(Stats, State = #state{gc_state = GcSt}) ->
case ?ENABLED(GcSt) andalso emqx_gc:run(Stats, GcSt) of
false -> State;
{IsGC, GcSt1} ->
IsGC andalso emqx_metrics:inc('channel.gc.cnt'),
State#state{gc_state = GcSt1}
end.
check_oom(State = #state{channel = Channel}) ->
#{zone := Zone} = emqx_channel:info(clientinfo, Channel),
OomPolicy = emqx_zone:oom_policy(Zone),
case ?ENABLED(OomPolicy) andalso emqx_misc:check_oom(OomPolicy) of
Shutdown = {shutdown, _Reason} ->
erlang:send(self(), Shutdown);
_Other -> ok
end,
State.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Activate Socket %% Activate Socket
@ -587,48 +660,30 @@ close_socket(State = #state{transport = Transport, socket = Socket}) ->
ok = Transport:fast_close(Socket), ok = Transport:fast_close(Socket),
State#state{sockstate = closed}. State#state{sockstate = closed}.
%%--------------------------------------------------------------------
%% Ensure rate limit
-define(ENABLED(Rl), (Rl =/= undefined)).
ensure_rate_limit(State = #state{rate_limit = Rl, pub_limit = Pl}) ->
Pubs = emqx_pd:reset_counter(incoming_pubs),
Bytes = emqx_pd:reset_counter(incoming_bytes),
Limiters = [{Pl, #state.pub_limit, Pubs} || ?ENABLED(Pl)] ++
[{Rl, #state.rate_limit, Bytes} || ?ENABLED(Rl)],
ensure_rate_limit(Limiters, State).
ensure_rate_limit([], State) ->
State;
ensure_rate_limit([{Rl, Pos, Cnt}|Limiters], State) ->
case esockd_rate_limit:check(Cnt, Rl) of
{0, Rl1} ->
ensure_rate_limit(Limiters, setelement(Pos, State, Rl1));
{Pause, Rl1} ->
?LOG(debug, "Pause ~pms due to rate limit", [Pause]),
TRef = erlang:send_after(Pause, self(), activate_socket),
NState = State#state{sockstate = blocked, limit_timer = TRef},
setelement(Pos, NState, Rl1)
end.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Inc incoming/outgoing stats %% Inc incoming/outgoing stats
-compile({inline, [inc_incoming_stats/1]}). -compile({inline, [inc_incoming_stats/1]}).
inc_incoming_stats(Type) when is_integer(Type) -> inc_incoming_stats(Packet = ?PACKET(Type)) ->
emqx_pd:update_counter(recv_pkt, 1), emqx_pd:inc_counter(recv_pkt, 1),
if if
Type == ?PUBLISH -> Type == ?PUBLISH ->
emqx_pd:update_counter(recv_msg, 1), emqx_pd:inc_counter(recv_msg, 1),
emqx_pd:update_counter(incoming_pubs, 1); emqx_pd:inc_counter(incoming_pubs, 1);
true -> ok true -> ok
end. end,
emqx_metrics:inc_recv(Packet).
-compile({inline, [inc_outgoing_stats/1]}). -compile({inline, [inc_outgoing_stats/1]}).
inc_outgoing_stats(Type) -> inc_outgoing_stats(Packet = ?PACKET(Type)) ->
emqx_pd:update_counter(send_pkt, 1), emqx_pd:inc_counter(send_pkt, 1),
(Type == ?PUBLISH) andalso emqx_pd:update_counter(send_msg, 1). if
Type == ?PUBLISH ->
emqx_pd:inc_counter(send_msg, 1),
emqx_pd:inc_counter(outgoing_pubs, 1);
true -> ok
end,
emqx_metrics:inc_sent(Packet).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Helper functions %% Helper functions
@ -646,13 +701,14 @@ next_msgs(Action) when is_tuple(Action) ->
next_msgs(Actions) when is_list(Actions) -> next_msgs(Actions) when is_list(Actions) ->
Actions. Actions.
-compile({inline, [shutdown/2, shutdown/3]}).
shutdown(Reason, State) -> shutdown(Reason, State) ->
stop({shutdown, Reason}, State). stop({shutdown, Reason}, State).
shutdown(Reason, Reply, State) -> shutdown(Reason, Reply, State) ->
stop({shutdown, Reason}, Reply, State). stop({shutdown, Reason}, Reply, State).
-compile({inline, [stop/2]}). -compile({inline, [stop/2, stop/3]}).
stop(Reason, State) -> stop(Reason, State) ->
{stop, Reason, State}. {stop, Reason, State}.

View File

@ -29,6 +29,7 @@
-include("types.hrl"). -include("types.hrl").
-export([ init/1 -export([ init/1
, run/2
, run/3 , run/3
, info/1 , info/1
, reset/1 , reset/1
@ -57,21 +58,26 @@ init(#{count := Count, bytes := Bytes}) ->
?GCS(maps:from_list(Cnt ++ Oct)). ?GCS(maps:from_list(Cnt ++ Oct)).
%% @doc Try to run GC based on reduntions of count or bytes. %% @doc Try to run GC based on reduntions of count or bytes.
-spec(run(#{cnt := pos_integer(), oct := pos_integer()}, gc_state())
-> {boolean(), gc_state()}).
run(#{cnt := Cnt, oct := Oct}, GcSt) ->
run(Cnt, Oct, GcSt).
-spec(run(pos_integer(), pos_integer(), gc_state()) -spec(run(pos_integer(), pos_integer(), gc_state())
-> {boolean(), gc_state()}). -> {boolean(), gc_state()}).
run(Cnt, Oct, ?GCS(St)) -> run(Cnt, Oct, ?GCS(St)) ->
{Res, St1} = run([{cnt, Cnt}, {oct, Oct}], St), {Res, St1} = do_run([{cnt, Cnt}, {oct, Oct}], St),
{Res, ?GCS(St1)}. {Res, ?GCS(St1)}.
run([], St) -> do_run([], St) ->
{false, St}; {false, St};
run([{K, N}|T], St) -> do_run([{K, N}|T], St) ->
case dec(K, N, St) of case dec(K, N, St) of
{true, St1} -> {true, St1} ->
true = erlang:garbage_collect(), erlang:garbage_collect(),
{true, do_reset(St1)}; {true, do_reset(St1)};
{false, St1} -> {false, St1} ->
run(T, St1) do_run(T, St1)
end. end.
%% @doc Info of GC state. %% @doc Info of GC state.

73
src/emqx_limiter.erl Normal file
View File

@ -0,0 +1,73 @@
%%--------------------------------------------------------------------
%% Copyright (c) 2019 EMQ Technologies Co., Ltd. All Rights Reserved.
%%
%% Licensed under the Apache License, Version 2.0 (the "License");
%% you may not use this file except in compliance with the License.
%% You may obtain a copy of the License at
%%
%% http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.
%%--------------------------------------------------------------------
-module(emqx_limiter).
-include("types.hrl").
-export([init/1, info/1, check/2]).
-import(emqx_misc, [maybe_apply/2]).
-record(limiter, {
%% Publish Limit
pub_limit :: maybe(esockd_rate_limit:bucket()),
%% Rate Limit
rate_limit :: maybe(esockd_rate_limit:bucket())
}).
-type(limiter() :: #limiter{}).
-export_type([limiter/0]).
-define(ENABLED(Rl), (Rl =/= undefined)).
-spec(init(proplists:proplist()) -> maybe(limiter())).
init(Options) ->
Zone = proplists:get_value(zone, Options),
Pl = emqx_zone:publish_limit(Zone),
Rl = proplists:get_value(rate_limit, Options),
case ?ENABLED(Pl) or ?ENABLED(Rl) of
true -> #limiter{pub_limit = init_limit(Pl),
rate_limit = init_limit(Rl)
};
false -> undefined
end.
init_limit(Rl) ->
maybe_apply(fun esockd_rate_limit:new/1, Rl).
info(#limiter{pub_limit = Pl, rate_limit = Rl}) ->
#{pub_limit => info(Pl), rate_limit => info(Rl)};
info(Rl) ->
maybe_apply(fun esockd_rate_limit:info/1, Rl).
check(#{cnt := Cnt, oct := Oct}, Limiter = #limiter{pub_limit = Pl,
rate_limit = Rl}) ->
do_check([{#limiter.pub_limit, Cnt, Pl} || ?ENABLED(Pl)] ++
[{#limiter.rate_limit, Oct, Rl} || ?ENABLED(Rl)], Limiter).
do_check([], Limiter) ->
{ok, Limiter};
do_check([{Pos, Cnt, Rl}|More], Limiter) ->
case esockd_rate_limit:check(Cnt, Rl) of
{0, Rl1} ->
do_check(More, setelement(Pos, Limiter, Rl1));
{Pause, Rl1} ->
{pause, Pause, setelement(Pos, Limiter, Rl1)}
end.

View File

@ -16,6 +16,8 @@
-module(emqx_misc). -module(emqx_misc).
-compile(inline).
-include("types.hrl"). -include("types.hrl").
-include("logger.hrl"). -include("logger.hrl").
@ -26,22 +28,20 @@
, start_timer/2 , start_timer/2
, start_timer/3 , start_timer/3
, cancel_timer/1 , cancel_timer/1
, drain_deliver/0
, drain_down/1
, check_oom/1
, check_oom/2
, tune_heap_size/1
, proc_name/2 , proc_name/2
, proc_stats/0 , proc_stats/0
, proc_stats/1 , proc_stats/1
, rand_seed/0
, now_to_secs/1
, now_to_ms/1
, index_of/2 , index_of/2
]). ]).
-export([ drain_deliver/0
, drain_deliver/1
, drain_down/1
]).
-compile({inline,
[ start_timer/2
, start_timer/3
]}).
%% @doc Merge options %% @doc Merge options
-spec(merge_opts(Opts, Opts) -> Opts when Opts :: proplists:proplist()). -spec(merge_opts(Opts, Opts) -> Opts when Opts :: proplists:proplist()).
merge_opts(Defaults, Options) -> merge_opts(Defaults, Options) ->
@ -112,6 +112,68 @@ cancel_timer(Timer) when is_reference(Timer) ->
end; end;
cancel_timer(_) -> ok. cancel_timer(_) -> ok.
%% @doc Drain delivers from the channel proc's mailbox.
drain_deliver() ->
drain_deliver([]).
drain_deliver(Acc) ->
receive
Deliver = {deliver, _Topic, _Msg} ->
drain_deliver([Deliver|Acc])
after 0 ->
lists:reverse(Acc)
end.
%% @doc Drain process 'DOWN' events.
-spec(drain_down(pos_integer()) -> list(pid())).
drain_down(Cnt) when Cnt > 0 ->
drain_down(Cnt, []).
drain_down(0, Acc) ->
lists:reverse(Acc);
drain_down(Cnt, Acc) ->
receive
{'DOWN', _MRef, process, Pid, _Reason} ->
drain_down(Cnt - 1, [Pid|Acc])
after 0 ->
lists:reverse(Acc)
end.
%% @doc Check process's mailbox and heapsize against OOM policy,
%% return `ok | {shutdown, Reason}' accordingly.
%% `ok': There is nothing out of the ordinary.
%% `shutdown': Some numbers (message queue length hit the limit),
%% hence shutdown for greater good (system stability).
-spec(check_oom(emqx_types:oom_policy()) -> ok | {shutdown, term()}).
check_oom(Policy) ->
check_oom(self(), Policy).
-spec(check_oom(pid(), emqx_types:oom_policy()) -> ok | {shutdown, term()}).
check_oom(Pid, #{message_queue_len := MaxQLen,
max_heap_size := MaxHeapSize}) ->
case process_info(Pid, [message_queue_len, total_heap_size]) of
undefined -> ok;
[{message_queue_len, QLen}, {total_heap_size, HeapSize}] ->
do_check_oom([{QLen, MaxQLen, message_queue_too_long},
{HeapSize, MaxHeapSize, proc_heap_too_large}
])
end.
do_check_oom([]) -> ok;
do_check_oom([{Val, Max, Reason}|Rest]) ->
case is_integer(Max) andalso (0 < Max) andalso (Max < Val) of
true -> {shutdown, Reason};
false -> do_check_oom(Rest)
end.
tune_heap_size(#{max_heap_size := MaxHeapSize}) ->
%% If set to zero, the limit is disabled.
erlang:process_flag(max_heap_size, #{size => MaxHeapSize,
kill => false,
error_logger => true
});
tune_heap_size(undefined) -> ok.
-spec(proc_name(atom(), pos_integer()) -> atom()). -spec(proc_name(atom(), pos_integer()) -> atom()).
proc_name(Mod, Id) -> proc_name(Mod, Id) ->
list_to_atom(lists:concat([Mod, "_", Id])). list_to_atom(lists:concat([Mod, "_", Id])).
@ -132,32 +194,16 @@ proc_stats(Pid) ->
[{mailbox_len, Len}|ProcStats] [{mailbox_len, Len}|ProcStats]
end. end.
%% @doc Drain delivers from the channel's mailbox. rand_seed() ->
drain_deliver() -> rand:seed(exsplus, erlang:timestamp()).
drain_deliver([]).
drain_deliver(Acc) -> -spec(now_to_secs(erlang:timestamp()) -> pos_integer()).
receive now_to_secs({MegaSecs, Secs, _MicroSecs}) ->
Deliver = {deliver, _Topic, _Msg} -> MegaSecs * 1000000 + Secs.
drain_deliver([Deliver|Acc])
after 0 ->
lists:reverse(Acc)
end.
%% @doc Drain process down events. -spec(now_to_ms(erlang:timestamp()) -> pos_integer()).
-spec(drain_down(pos_integer()) -> list(pid())). now_to_ms({MegaSecs, Secs, MicroSecs}) ->
drain_down(Cnt) when Cnt > 0 -> (MegaSecs * 1000000 + Secs) * 1000 + round(MicroSecs/1000).
drain_down(Cnt, []).
drain_down(0, Acc) ->
lists:reverse(Acc);
drain_down(Cnt, Acc) ->
receive
{'DOWN', _MRef, process, Pid, _Reason} ->
drain_down(Cnt - 1, [Pid|Acc])
after 0 ->
drain_down(0, Acc)
end.
%% lists:index_of/2 %% lists:index_of/2
index_of(E, L) -> index_of(E, L) ->

View File

@ -21,14 +21,14 @@
-export([ get_counters/1 -export([ get_counters/1
, get_counter/1 , get_counter/1
, update_counter/2 , inc_counter/2
, reset_counter/1 , reset_counter/1
]). ]).
-compile({inline, -compile({inline,
[ get_counters/1 [ get_counters/1
, get_counter/1 , get_counter/1
, update_counter/2 , inc_counter/2
, reset_counter/1 , reset_counter/1
]}). ]}).
@ -42,8 +42,8 @@ get_counters(Keys) when is_list(Keys) ->
get_counter(Key) -> get_counter(Key) ->
case get(Key) of undefined -> 0; Cnt -> Cnt end. case get(Key) of undefined -> 0; Cnt -> Cnt end.
-spec(update_counter(key(), number()) -> maybe(number())). -spec(inc_counter(key(), number()) -> maybe(number())).
update_counter(Key, Inc) -> inc_counter(Key, Inc) ->
put(Key, get_counter(Key) + Inc). put(Key, get_counter(Key) + Inc).
-spec(reset_counter(key()) -> number()). -spec(reset_counter(key()) -> number()).

View File

@ -85,6 +85,8 @@
, stats/0 , stats/0
]). ]).
-export_type([oom_policy/0]).
-type(ver() :: ?MQTT_PROTO_V3 -type(ver() :: ?MQTT_PROTO_V3
| ?MQTT_PROTO_V4 | ?MQTT_PROTO_V4
| ?MQTT_PROTO_V5). | ?MQTT_PROTO_V5).
@ -186,3 +188,7 @@
-type(infos() :: #{atom() => term()}). -type(infos() :: #{atom() => term()}).
-type(stats() :: #{atom() => non_neg_integer()|stats()}). -type(stats() :: #{atom() => non_neg_integer()|stats()}).
-type(oom_policy() :: #{message_queue_len => non_neg_integer(),
max_heap_size => non_neg_integer()
}).

View File

@ -45,6 +45,11 @@
, terminate/3 , terminate/3
]). ]).
-import(emqx_misc,
[ maybe_apply/2
, start_timer/2
]).
-record(state, { -record(state, {
%% Peername of the ws connection. %% Peername of the ws connection.
peername :: emqx_types:peername(), peername :: emqx_types:peername(),
@ -52,24 +57,41 @@
sockname :: emqx_types:peername(), sockname :: emqx_types:peername(),
%% Sock state %% Sock state
sockstate :: emqx_types:sockstate(), sockstate :: emqx_types:sockstate(),
%% Parser State %% Simulate the active_n opt
active_n :: pos_integer(),
%% Limiter
limiter :: emqx_limiter:limiter(),
%% Limit Timer
limit_timer :: maybe(reference()),
%% Parse State
parse_state :: emqx_frame:parse_state(), parse_state :: emqx_frame:parse_state(),
%% Serialize function %% Serialize function
serialize :: emqx_frame:serialize_fun(), serialize :: emqx_frame:serialize_fun(),
%% Channel %% Channel
channel :: emqx_channel:channel(), channel :: emqx_channel:channel(),
%% GC State
gc_state :: maybe(emqx_gc:gc_state()),
%% Out Pending Packets %% Out Pending Packets
pendings :: list(emqx_types:packet()), pendings :: list(emqx_types:packet()),
%% Stats Timer
stats_timer :: disabled | maybe(reference()),
%% Idle Timeout
idle_timeout :: timeout(),
%% Idle Timer
idle_timer :: reference(),
%% The stop reason %% The stop reason
stop_reason :: term() stop_reason :: term()
}). }).
-type(state() :: #state{}). -type(state() :: #state{}).
-define(INFO_KEYS, [socktype, peername, sockname, sockstate]). -define(ACTIVE_N, 100).
-define(INFO_KEYS, [socktype, peername, sockname, sockstate, active_n, limiter]).
-define(SOCK_STATS, [recv_oct, recv_cnt, send_oct, send_cnt]). -define(SOCK_STATS, [recv_oct, recv_cnt, send_oct, send_cnt]).
-define(CONN_STATS, [recv_pkt, recv_msg, send_pkt, send_msg]). -define(CONN_STATS, [recv_pkt, recv_msg, send_pkt, send_msg]).
-define(ENABLED(X), (X =/= undefined)).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% API %% API
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
@ -92,11 +114,20 @@ info(sockname, #state{sockname = Sockname}) ->
Sockname; Sockname;
info(sockstate, #state{sockstate = SockSt}) -> info(sockstate, #state{sockstate = SockSt}) ->
SockSt; SockSt;
info(active_n, #state{active_n = ActiveN}) ->
ActiveN;
info(limiter, #state{limiter = Limiter}) ->
maybe_apply(fun emqx_limiter:info/1, Limiter);
info(channel, #state{channel = Channel}) -> info(channel, #state{channel = Channel}) ->
emqx_channel:info(Channel); emqx_channel:info(Channel);
info(stop_reason, #state{stop_reason = Reason}) -> info(stop_reason, #state{stop_reason = Reason}) ->
Reason. Reason.
attrs(State = #state{channel = Channel}) ->
ChanAttrs = emqx_channel:attrs(Channel),
SockAttrs = maps:from_list(info(?INFO_KEYS, State)),
maps:merge(ChanAttrs, #{sockinfo => SockAttrs}).
-spec(stats(pid()|state()) -> emqx_types:stats()). -spec(stats(pid()|state()) -> emqx_types:stats()).
stats(WsPid) when is_pid(WsPid) -> stats(WsPid) when is_pid(WsPid) ->
call(WsPid, stats); call(WsPid, stats);
@ -128,6 +159,7 @@ call(WsPid, Req) when is_pid(WsPid) ->
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
init(Req, Opts) -> init(Req, Opts) ->
%% WS Transport Idle Timeout
IdleTimeout = proplists:get_value(idle_timeout, Opts, 7200000), IdleTimeout = proplists:get_value(idle_timeout, Opts, 7200000),
DeflateOptions = maps:from_list(proplists:get_value(deflate_options, Opts, [])), DeflateOptions = maps:from_list(proplists:get_value(deflate_options, Opts, [])),
MaxFrameSize = case proplists:get_value(max_frame_size, Opts, 0) of MaxFrameSize = case proplists:get_value(max_frame_size, Opts, 0) of
@ -174,29 +206,41 @@ websocket_init([Req, Opts]) ->
conn_mod => ?MODULE conn_mod => ?MODULE
}, },
Zone = proplists:get_value(zone, Opts), Zone = proplists:get_value(zone, Opts),
FrameOpts = emqx_zone:frame_options(Zone), ActiveN = proplists:get_value(active_n, Opts, ?ACTIVE_N),
Limiter = emqx_limiter:init(Opts),
FrameOpts = emqx_zone:mqtt_frame_options(Zone),
ParseState = emqx_frame:initial_parse_state(FrameOpts), ParseState = emqx_frame:initial_parse_state(FrameOpts),
Serialize = emqx_frame:serialize_fun(), Serialize = emqx_frame:serialize_fun(),
Channel = emqx_channel:init(ConnInfo, Opts), Channel = emqx_channel:init(ConnInfo, Opts),
GcState = emqx_zone:init_gc_state(Zone),
StatsTimer = emqx_zone:stats_timer(Zone),
%% MQTT Idle Timeout
IdleTimeout = emqx_zone:idle_timeout(Zone),
IdleTimer = start_timer(IdleTimeout, idle_timeout),
emqx_misc:tune_heap_size(emqx_zone:oom_policy(Zone)),
emqx_logger:set_metadata_peername(esockd_net:format(Peername)), emqx_logger:set_metadata_peername(esockd_net:format(Peername)),
{ok, #state{peername = Peername, {ok, #state{peername = Peername,
sockname = Sockname, sockname = Sockname,
sockstate = idle, sockstate = running,
parse_state = ParseState, active_n = ActiveN,
serialize = Serialize, limiter = Limiter,
channel = Channel, parse_state = ParseState,
pendings = [] serialize = Serialize,
}}. channel = Channel,
gc_state = GcState,
pendings = [],
stats_timer = StatsTimer,
idle_timeout = IdleTimeout,
idle_timer = IdleTimer
}, hibernate}.
websocket_handle({binary, Data}, State) when is_list(Data) -> websocket_handle({binary, Data}, State) when is_list(Data) ->
websocket_handle({binary, iolist_to_binary(Data)}, State); websocket_handle({binary, iolist_to_binary(Data)}, State);
websocket_handle({binary, Data}, State = #state{channel = Channel}) -> websocket_handle({binary, Data}, State) ->
?LOG(debug, "RECV ~p", [Data]), ?LOG(debug, "RECV ~p", [Data]),
Oct = iolist_size(Data), ok = inc_recv_stats(1, iolist_size(Data)),
ok = inc_recv_stats(1, Oct), parse_incoming(Data, ensure_stats_timer(State));
NChannel = emqx_channel:recvd(Oct, Channel),
parse_incoming(Data, State#state{channel = NChannel});
%% Pings should be replied with pongs, cowboy does it automatically %% Pings should be replied with pongs, cowboy does it automatically
%% Pongs can be safely ignored. Clause here simply prevents crash. %% Pongs can be safely ignored. Clause here simply prevents crash.
@ -215,30 +259,43 @@ websocket_info({call, From, Req}, State) ->
handle_call(From, Req, State); handle_call(From, Req, State);
websocket_info({cast, Msg}, State = #state{channel = Channel}) -> websocket_info({cast, Msg}, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_info(Msg, Channel), State); handle_chan_return(emqx_channel:handle_info(Msg, Channel), State);
websocket_info({incoming, Packet = ?CONNECT_PACKET(ConnPkt)}, State) -> websocket_info({incoming, Packet = ?CONNECT_PACKET(ConnPkt)},
State = #state{idle_timer = IdleTimer}) ->
ok = emqx_misc:cancel_timer(IdleTimer),
Serialize = emqx_frame:serialize_fun(ConnPkt), Serialize = emqx_frame:serialize_fun(ConnPkt),
NState = State#state{sockstate = running, NState = State#state{serialize = Serialize,
serialize = Serialize idle_timer = undefined
}, },
handle_incoming(Packet, NState); handle_incoming(Packet, NState);
websocket_info({incoming, ?PACKET(?PINGREQ)}, State) ->
reply(?PACKET(?PINGRESP), State);
websocket_info({incoming, Packet}, State) -> websocket_info({incoming, Packet}, State) ->
handle_incoming(Packet, State); handle_incoming(Packet, State);
websocket_info(Deliver = {deliver, _Topic, _Msg}, websocket_info(rate_limit, State) ->
State = #state{channel = Channel}) -> InStats = #{cnt => emqx_pd:reset_counter(incoming_pubs),
Delivers = emqx_misc:drain_deliver([Deliver]), oct => emqx_pd:reset_counter(incoming_bytes)
Result = emqx_channel:handle_out(Delivers, Channel), },
handle_return(Result, State); erlang:send(self(), {check_gc, InStats}),
ensure_rate_limit(InStats, State);
websocket_info({timeout, TRef, keepalive}, State) when is_reference(TRef) -> websocket_info({check_gc, Stats}, State) ->
RecvOct = emqx_pd:get_counter(recv_oct), {ok, check_oom(run_gc(Stats, State))};
handle_timeout(TRef, {keepalive, RecvOct}, State);
websocket_info({timeout, TRef, emit_stats}, State) when is_reference(TRef) -> websocket_info({deliver, _Topic, _Msg} = Deliver, State = #state{channel = Channel}) ->
handle_timeout(TRef, {emit_stats, stats(State)}, State); Delivers = [Deliver|emqx_misc:drain_deliver()],
Ret = emqx_channel:handle_out(Delivers, Channel),
handle_chan_return(Ret, State);
websocket_info({timeout, TRef, limit_timeout}, State = #state{limit_timer = TRef}) ->
NState = State#state{sockstate = running,
limit_timer = undefined
},
{reply, [{active, true}], NState};
websocket_info({timeout, TRef, Msg}, State) when is_reference(TRef) -> websocket_info({timeout, TRef, Msg}, State) when is_reference(TRef) ->
handle_timeout(TRef, Msg, State); handle_timeout(TRef, Msg, State);
@ -293,27 +350,89 @@ handle_call(From, Req, State = #state{channel = Channel}) ->
%% Handle Info %% Handle Info
handle_info({connack, ConnAck}, State = #state{channel = Channel}) -> handle_info({connack, ConnAck}, State = #state{channel = Channel}) ->
ChanAttrs = emqx_channel:attrs(Channel), ClientId = emqx_channel:info(clientid, Channel),
SockAttrs = maps:from_list(info(?INFO_KEYS, State)), ok = emqx_cm:register_channel(ClientId),
Attrs = maps:merge(ChanAttrs, #{sockinfo => SockAttrs}), ok = emqx_cm:set_chan_attrs(ClientId, attrs(State)),
ok = emqx_channel:handle_info({register, Attrs, stats(State)}, Channel), ok = emqx_cm:set_chan_stats(ClientId, stats(State)),
reply(enqueue(ConnAck, State)); reply(enqueue(ConnAck, State));
handle_info({enter, disconnected}, State = #state{channel = Channel}) -> handle_info({enter, disconnected}, State = #state{channel = Channel}) ->
ChanAttrs = emqx_channel:attrs(Channel), ClientId = emqx_channel:info(clientid, Channel),
SockAttrs = maps:from_list(info(?INFO_KEYS, State)), emqx_cm:set_chan_attrs(ClientId, attrs(State)),
Attrs = maps:merge(ChanAttrs, #{sockinfo => SockAttrs}), emqx_cm:set_chan_stats(ClientId, stats(State)),
ok = emqx_channel:handle_info({register, Attrs, stats(State)}, Channel),
reply(State); reply(State);
handle_info(Info, State = #state{channel = Channel}) -> handle_info(Info, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_info(Info, Channel), State). Ret = emqx_channel:handle_info(Info, Channel),
handle_chan_return(Ret, State).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle timeout %% Handle timeout
handle_timeout(TRef, Msg, State = #state{channel = Channel}) -> handle_timeout(TRef, idle_timeout, State = #state{idle_timer = TRef}) ->
handle_return(emqx_channel:handle_timeout(TRef, Msg, Channel), State). shutdown(idle_timeout, State);
handle_timeout(TRef, keepalive, State) when is_reference(TRef) ->
RecvOct = emqx_pd:get_counter(recv_oct),
handle_timeout(TRef, {keepalive, RecvOct}, State);
handle_timeout(TRef, emit_stats, State = #state{channel = Channel,
stats_timer = TRef}) ->
ClientId = emqx_channel:info(clientid, Channel),
(ClientId =/= undefined) andalso emqx_cm:set_chan_stats(ClientId, stats(State)),
reply(State#state{stats_timer = undefined});
handle_timeout(TRef, TMsg, State = #state{channel = Channel}) ->
Ret = emqx_channel:handle_timeout(TRef, TMsg, Channel),
handle_chan_return(Ret, State).
%%--------------------------------------------------------------------
%% Ensure stats timer
-compile({inline, [ensure_stats_timer/1]}).
ensure_stats_timer(State = #state{idle_timeout = Timeout,
stats_timer = undefined}) ->
State#state{stats_timer = start_timer(Timeout, emit_stats)};
ensure_stats_timer(State) -> State.
%%--------------------------------------------------------------------
%% Ensure rate limit
ensure_rate_limit(Stats, State = #state{limiter = Limiter}) ->
case ?ENABLED(Limiter) andalso emqx_limiter:check(Stats, Limiter) of
false -> {ok, State};
{ok, Limiter1} ->
{ok, State#state{limiter = Limiter1}};
{pause, Time, Limiter1} ->
?LOG(debug, "Pause ~pms due to rate limit", [Time]),
TRef = start_timer(Time, limit_timeout),
NState = State#state{sockstate = blocked,
limiter = Limiter1,
limit_timer = TRef
},
{reply, [{active, false}], NState}
end.
%%--------------------------------------------------------------------
%% Run GC and Check OOM
run_gc(Stats, State = #state{gc_state = GcSt}) ->
case ?ENABLED(GcSt) andalso emqx_gc:run(Stats, GcSt) of
false -> State;
{IsGC, GcSt1} ->
IsGC andalso emqx_metrics:inc('channel.gc.cnt'),
State#state{gc_state = GcSt1}
end.
check_oom(State = #state{channel = Channel}) ->
#{zone := Zone} = emqx_channel:info(clientinfo, Channel),
OomPolicy = emqx_zone:oom_policy(Zone),
case ?ENABLED(OomPolicy) andalso emqx_misc:check_oom(OomPolicy) of
Shutdown = {shutdown, _Reason} ->
erlang:send(self(), Shutdown);
_Other -> ok
end,
State.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Parse incoming data %% Parse incoming data
@ -326,7 +445,7 @@ parse_incoming(Data, State = #state{parse_state = ParseState}) ->
{more, NParseState} -> {more, NParseState} ->
{ok, State#state{parse_state = NParseState}}; {ok, State#state{parse_state = NParseState}};
{ok, Packet, Rest, NParseState} -> {ok, Packet, Rest, NParseState} ->
self() ! {incoming, Packet}, erlang:send(self(), {incoming, Packet}),
parse_incoming(Rest, State#state{parse_state = NParseState}) parse_incoming(Rest, State#state{parse_state = NParseState})
catch catch
error:Reason:Stk -> error:Reason:Stk ->
@ -337,52 +456,60 @@ parse_incoming(Data, State = #state{parse_state = ParseState}) ->
end. end.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle incoming packets %% Handle incoming packet
handle_incoming(Packet = ?PACKET(Type), State = #state{channel = Channel}) -> handle_incoming(Packet, State = #state{active_n = ActiveN, channel = Channel})
_ = inc_incoming_stats(Type), when is_record(Packet, mqtt_packet) ->
_ = emqx_metrics:inc_recv(Packet),
?LOG(debug, "RECV ~s", [emqx_packet:format(Packet)]), ?LOG(debug, "RECV ~s", [emqx_packet:format(Packet)]),
handle_return(emqx_channel:handle_in(Packet, Channel), State); ok = inc_incoming_stats(Packet),
(emqx_pd:get_counter(incoming_pubs) > ActiveN)
andalso erlang:send(self(), rate_limit),
Ret = emqx_channel:handle_in(Packet, Channel),
handle_chan_return(Ret, State);
handle_incoming(FrameError, State = #state{channel = Channel}) -> handle_incoming(FrameError, State = #state{channel = Channel}) ->
handle_return(emqx_channel:handle_in(FrameError, Channel), State). handle_chan_return(emqx_channel:handle_in(FrameError, Channel), State).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle channel return %% Handle channel return
handle_return(ok, State) -> handle_chan_return(ok, State) ->
reply(State); reply(State);
handle_return({ok, NChannel}, State) -> handle_chan_return({ok, NChannel}, State) ->
reply(State#state{channel= NChannel}); reply(State#state{channel= NChannel});
handle_return({ok, Replies, NChannel}, State) -> handle_chan_return({ok, Replies, NChannel}, State) ->
reply(Replies, State#state{channel= NChannel}); reply(Replies, State#state{channel= NChannel});
handle_return({shutdown, Reason, NChannel}, State) -> handle_chan_return({shutdown, Reason, NChannel}, State) ->
stop(Reason, State#state{channel = NChannel}); stop(Reason, State#state{channel = NChannel});
handle_return({shutdown, Reason, OutPacket, NChannel}, State) -> handle_chan_return({shutdown, Reason, OutPacket, NChannel}, State) ->
NState = State#state{channel = NChannel}, NState = State#state{channel = NChannel},
stop(Reason, enqueue(OutPacket, NState)). stop(Reason, enqueue(OutPacket, NState)).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% Handle outgoing packets %% Handle outgoing packets
handle_outgoing(Packets, State = #state{channel = Channel}) -> handle_outgoing(Packets, State = #state{active_n = ActiveN}) ->
IoData = lists:map(serialize_and_inc_stats_fun(State), Packets), IoData = lists:map(serialize_and_inc_stats_fun(State), Packets),
Oct = iolist_size(IoData), Oct = iolist_size(IoData),
ok = inc_sent_stats(length(Packets), Oct), ok = inc_sent_stats(length(Packets), Oct),
NChannel = emqx_channel:sent(Oct, Channel), case emqx_pd:get_counter(outgoing_pubs) > ActiveN of
{{binary, IoData}, State#state{channel = NChannel}}. true ->
OutStats = #{cnt => emqx_pd:reset_counter(outgoing_pubs),
oct => emqx_pd:reset_counter(outgoing_bytes)
},
erlang:send(self(), {check_gc, OutStats});
false -> ok
end,
{{binary, IoData}, ensure_stats_timer(State)}.
%% TODO: Duplicated with emqx_channel:serialize_and_inc_stats_fun/1
serialize_and_inc_stats_fun(#state{serialize = Serialize}) -> serialize_and_inc_stats_fun(#state{serialize = Serialize}) ->
fun(Packet = ?PACKET(Type)) -> fun(Packet) ->
case Serialize(Packet) of case Serialize(Packet) of
<<>> -> ?LOG(warning, "~s is discarded due to the frame is too large!", <<>> -> ?LOG(warning, "~s is discarded due to the frame is too large!",
[emqx_packet:format(Packet)]), [emqx_packet:format(Packet)]),
<<>>; <<>>;
Data -> _ = inc_outgoing_stats(Type), Data -> ?LOG(debug, "SEND ~s", [emqx_packet:format(Packet)]),
_ = emqx_metrics:inc_sent(Packet), ok = inc_outgoing_stats(Packet),
?LOG(debug, "SEND ~s", [emqx_packet:format(Packet)]),
Data Data
end end
end. end.
@ -398,23 +525,33 @@ serialize_and_inc_stats_fun(#state{serialize = Serialize}) ->
]}). ]}).
inc_recv_stats(Cnt, Oct) -> inc_recv_stats(Cnt, Oct) ->
emqx_pd:update_counter(recv_cnt, Cnt), emqx_pd:inc_counter(incoming_bytes, Oct),
emqx_pd:update_counter(recv_oct, Oct), emqx_pd:inc_counter(recv_cnt, Cnt),
emqx_pd:inc_counter(recv_oct, Oct),
emqx_metrics:inc('bytes.received', Oct). emqx_metrics:inc('bytes.received', Oct).
inc_incoming_stats(Type) -> inc_incoming_stats(Packet = ?PACKET(Type)) ->
emqx_pd:update_counter(recv_pkt, 1), emqx_pd:inc_counter(recv_pkt, 1),
(Type == ?PUBLISH) if Type == ?PUBLISH ->
andalso emqx_pd:update_counter(recv_msg, 1). emqx_pd:inc_counter(recv_msg, 1),
emqx_pd:inc_counter(incoming_pubs, 1);
true -> ok
end,
emqx_metrics:inc_recv(Packet).
inc_outgoing_stats(Type) -> inc_outgoing_stats(Packet = ?PACKET(Type)) ->
emqx_pd:update_counter(send_pkt, 1), emqx_pd:inc_counter(send_pkt, 1),
(Type == ?PUBLISH) if Type == ?PUBLISH ->
andalso emqx_pd:update_counter(send_msg, 1). emqx_pd:inc_counter(send_msg, 1),
emqx_pd:inc_counter(outgoing_pubs, 1);
true -> ok
end,
emqx_metrics:inc_sent(Packet).
inc_sent_stats(Cnt, Oct) -> inc_sent_stats(Cnt, Oct) ->
emqx_pd:update_counter(send_cnt, Cnt), emqx_pd:inc_counter(outgoing_bytes, Oct),
emqx_pd:update_counter(send_oct, Oct), emqx_pd:inc_counter(send_cnt, Cnt),
emqx_pd:inc_counter(send_oct, Oct),
emqx_metrics:inc('bytes.sent', Oct). emqx_metrics:inc('bytes.sent', Oct).
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
@ -451,6 +588,9 @@ enqueue(Packet, State) when is_record(Packet, mqtt_packet) ->
enqueue(Packets, State = #state{pendings = Pendings}) -> enqueue(Packets, State = #state{pendings = Pendings}) ->
State#state{pendings = lists:append(Pendings, Packets)}. State#state{pendings = lists:append(Pendings, Packets)}.
shutdown(Reason, State) ->
stop({shutdown, Reason}, State).
stop(Reason, State = #state{pendings = []}) -> stop(Reason, State = #state{pendings = []}) ->
{stop, State#state{stop_reason = Reason}}; {stop, State#state{stop_reason = Reason}};
stop(Reason, State = #state{pendings = Pendings}) -> stop(Reason, State = #state{pendings = Pendings}) ->