port_command

src/emqttd_client.erl

@@ -20,7 +20,7 @@
 %%% SOFTWARE.
 %%%-----------------------------------------------------------------------------
 %%% @doc
-%%% MQTT Client
+%%% MQTT Client Connection.
 %%%
 %%% @end
 %%%-----------------------------------------------------------------------------
@@ -52,7 +52,7 @@
                 conn_name,
                 await_recv,
                 conn_state,
-                conserve,
+                rate_limiter,
                 parser,
                 proto_state,
                 packet_opts,
@@ -86,18 +86,22 @@ unsubscribe(CPid, Topics) ->
 init([SockArgs = {Transport, Sock, _SockFun}, MqttEnv]) ->
     % Transform if ssl.
     {ok, NewSock}  = esockd_connection:accept(SockArgs),
+    %%TODO:...
+    {ok, BufSizes} = inet:getopts(Sock, [sndbuf, recbuf, buffer]),
+    io:format("~p~n", [BufSizes]),
     {ok, Peername} = emqttd_net:peername(Sock),
     {ok, ConnStr}  = emqttd_net:connection_string(Sock, inbound),
-    SendFun = fun(Data) -> Transport:send(NewSock, Data) end,
+    SendFun    = send_fun(Transport, NewSock),
     PktOpts    = proplists:get_value(packet, MqttEnv),
     ProtoState = emqttd_protocol:init(Peername, SendFun, PktOpts),
-    State = control_throttle(#state{transport    = Transport,
+    Limiter    = proplists:get_value(rate_limiter, MqttEnv),
+    State      = run_socket(#state{transport    = Transport,
                                    socket       = NewSock,
                                    peername     = Peername,
                                    conn_name    = ConnStr,
                                    await_recv   = false,
                                    conn_state   = running,
-                                    conserve     = false,
+                                   rate_limiter = Limiter,
                                    packet_opts  = PktOpts,
                                    parser       = emqttd_parser:new(PktOpts),
                                    proto_state  = ProtoState}),
@@ -146,20 +150,26 @@ handle_info({redeliver, {?PUBREL, PacketId}},  State = #state{proto_state = Prot
     {ok, ProtoState1} = emqttd_protocol:redeliver({?PUBREL, PacketId}, ProtoState),
     noreply(State#state{proto_state = ProtoState1});
 
-handle_info({inet_reply, _Ref, ok}, State) ->
+handle_info(activate_sock, State) ->
-    noreply(State);
+    noreply(run_socket(State#state{conn_state = running}));
 
 handle_info({inet_async, Sock, _Ref, {ok, Data}}, State = #state{peername = Peername, socket = Sock}) ->
+    Size = size(Data),
     lager:debug("RECV from ~s: ~p", [emqttd_net:format(Peername), Data]),
-    emqttd_metrics:inc('bytes/received', size(Data)),
+    emqttd_metrics:inc('bytes/received', Size),
-    received(Data, control_throttle(State #state{await_recv = false}));
+    received(Data, rate_limit(Size, State#state{await_recv = false}));
 
 handle_info({inet_async, _Sock, _Ref, {error, Reason}}, State) ->
+    %%TODO: ...
     network_error(Reason, State);
 
+handle_info({inet_reply, _Ref, ok}, State) ->
+    %%TODO: ok...
+    io:format("inet_reply ok~n"),
+    noreply(State);
+
 handle_info({inet_reply, _Sock, {error, Reason}}, State) ->
-    ?ERROR("Unexpected inet_reply - ~p", [Reason], State),
+    network_error(Reason, State);
-    {noreply, State};
 
 handle_info({keepalive, start, TimeoutSec}, State = #state{transport = Transport, socket = Socket}) ->
     ?DEBUG("Start KeepAlive with ~p seconds", [TimeoutSec], State),
@@ -223,14 +233,14 @@ with_session(Fun, State = #state{proto_state = ProtoState}) ->
 
 %% receive and parse tcp data
 received(<<>>, State) ->
-    {noreply, State, hibernate};
+    noreply(State);
 
-received(Bytes, State = #state{packet_opts = PacketOpts,
+received(Bytes, State = #state{parser      = Parser,
-                               parser      = Parser,
+                               packet_opts = PacketOpts,
                                proto_state = ProtoState}) ->
     case catch Parser(Bytes) of
         {more, NewParser} ->
-            noreply(control_throttle(State#state{parser = NewParser}));
+            noreply(run_socket(State#state{parser = NewParser}));
         {ok, Packet, Rest} ->
             emqttd_metrics:received(Packet),
             case emqttd_protocol:received(Packet, ProtoState) of
@@ -258,6 +268,20 @@ network_error(Reason, State = #state{peername = Peername}) ->
                       [emqttd_net:format(Peername), Reason]),
     stop({shutdown, conn_closed}, State).
 
+rate_limit(_Size, State = #state{rate_limiter = undefined}) ->
+    run_socket(State);
+rate_limit(Size, State = #state{socket = Sock, rate_limiter = Limiter}) ->
+    {ok, BufSizes} = inet:getopts(Sock, [sndbuf, recbuf, buffer]),
+    io:format("~p~n", [BufSizes]),
+    case esockd_rate_limiter:check(Limiter, Size) of
+        {0, Limiter1} ->
+            run_socket(State#state{conn_state = running, rate_limiter = Limiter1});
+        {Pause, Limiter1} ->
+            ?ERROR("~p Received, Rate Limiter Pause for ~w", [Size, Pause], State),
+            erlang:send_after(Pause, self(), activate_sock),
+            State#state{conn_state = blocked, rate_limiter = Limiter1}    
+    end.
+
 run_socket(State = #state{conn_state = blocked}) ->
     State;
 run_socket(State = #state{await_recv = true}) ->
@@ -266,11 +290,11 @@ run_socket(State = #state{transport = Transport, socket = Sock}) ->
     Transport:async_recv(Sock, 0, infinity),
     State#state{await_recv = true}.
 
-control_throttle(State = #state{conn_state = Flow,
+send_fun(Transport, Sock) ->
-                                conserve   = Conserve}) ->
+    fun(Data) ->
-    case {Flow, Conserve} of
+        try Transport:port_command(Sock, Data) of
-        {running,   true} -> State #state{conn_state = blocked};
+            true -> ok
-        {blocked,  false} -> run_socket(State #state{conn_state = running});
+        catch
-        {_,            _} -> run_socket(State)
+            error:Error -> exit({socket_error, Error})
+        end
     end.
-