Add support for hearbeats

essen · essen · commit 9c153b2d4054 · 2016-01-06T17:54:11.000+01:00
This involves creating the keepalive supervisor when the connection
starts, and passing its pid onward to the handler. To achieve this
we put ourselves in just before starting cowboy_protocol, create
the supervisor and put its pid in the Cowboy environment. Later we
have a middleware that takes this pid and pass it to the MQTT handler
where it can then be used.

A few minor changes and fixes are also included in this commit,
including compacting the Req object and using hibernate.
diff --git a/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_app.erl b/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_app.erl
@@ -41,23 +41,28 @@ init([]) -> {ok, {{one_for_one, 1, 5}, []}}.
 
 mqtt_init() ->
     NbAcceptors = get_env(nb_acceptors, 1),
-    CowboyOpts = get_env(cowboy_opts, []),
+    CowboyOpts0 = get_env(cowboy_opts, []),
 
     Routes = cowboy_router:compile([{'_', [
         {"/example", cowboy_static, {priv_file, rabbitmq_web_mqtt, "example/index.html"}},
         {"/example/[...]", cowboy_static, {priv_dir, rabbitmq_web_mqtt, "example/"}},
         {"/ws", rabbit_web_mqtt_handler, []}
     ]}]),
+    CowboyOpts = [
+        {env, [{dispatch, Routes}]},
+        {middlewares, [cowboy_router, rabbit_web_mqtt_middleware, cowboy_handler]}
+    |CowboyOpts0],
 
-    TCPConf0 = get_env(tcp_config, []),
+    TCPConf0 = [{connection_type, supervisor}|get_env(tcp_config, [])],
     TCPConf = case proplists:get_value(port, TCPConf0) of
         undefined -> [{port, 15675}|TCPConf0];
         _ -> TCPConf0
     end,
     TCPPort = proplists:get_value(port, TCPConf),
 
-    {ok, _} = cowboy:start_http(web_mqtt, NbAcceptors, TCPConf,
-                                [{env, [{dispatch, Routes}]}|CowboyOpts]),
+    {ok, _} = ranch:start_listener(web_mqtt, NbAcceptors,
+        ranch_tcp, TCPConf,
+        rabbit_web_mqtt_connection_sup, CowboyOpts),
 
     rabbit_log:info("rabbit_web_mqtt: listening for HTTP connections on ~s:~w~n",
                     ["0.0.0.0", TCPPort]),
@@ -69,8 +74,9 @@ mqtt_init() ->
             rabbit_networking:ensure_ssl(),
             SSLPort = proplists:get_value(port, SSLConf),
 
-            {ok, _} = cowboy:start_https(web_mqtt_secure, NbAcceptors, SSLPort,
-                                         [{env, [{dispatch, Routes}]}|CowboyOpts]),
+            {ok, _} = cowboy:start_https(web_mqtt_secure, NbAcceptors,
+                ranch_ssl, SSLConf,
+                rabbit_web_mqtt_connection_sup, CowboyOpts),
             rabbit_log:info("rabbit_web_mqtt: listening for HTTPS connections on ~s:~w~n",
                             ["0.0.0.0", SSLPort])
     end,
diff --git a/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_connection_sup.erl b/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_connection_sup.erl
@@ -0,0 +1,59 @@
+%% The contents of this file are subject to the Mozilla Public License
+%% Version 1.1 (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.mozilla.org/MPL/
+%%
+%% Software distributed under the License is distributed on an "AS IS"
+%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
+%% the License for the specific language governing rights and
+%% limitations under the License.
+%%
+%% The Original Code is RabbitMQ.
+%%
+%% The Initial Developer of the Original Code is GoPivotal, Inc.
+%% Copyright (c) 2016 Pivotal Software, Inc.  All rights reserved.
+%%
+
+-module(rabbit_web_mqtt_connection_sup).
+
+-behaviour(supervisor2).
+-behaviour(ranch_protocol).
+
+-define(MAX_WAIT, 16#ffffffff).
+
+-export([start_link/4, start_keepalive_link/0]).
+
+-export([init/1]).
+
+%%----------------------------------------------------------------------------
+
+start_link(Ref, Sock, Transport, CowboyOpts0) ->
+    {ok, SupPid} = supervisor2:start_link(?MODULE, []),
+    {ok, KeepaliveSup} = supervisor2:start_child(
+                          SupPid,
+                          {rabbit_web_mqtt_keepalive_sup,
+                           {?MODULE, start_keepalive_link, []},
+                           intrinsic, infinity, supervisor, [rabbit_keepalive_sup]}),
+
+    %% In order for the Websocket handler to receive the KeepaliveSup
+    %% variable, we need to pass it first through the environment and
+    %% then have the middleware rabbit_web_mqtt_middleware place it
+    %% in the initial handler state.
+    {env, Env} = lists:keyfind(env, 1, CowboyOpts0),
+    CowboyOpts = lists:keyreplace(env, 1, CowboyOpts0,
+        {env, [{keepalive_sup, KeepaliveSup}|Env]}),
+
+    {ok, ReaderPid} = supervisor2:start_child(
+                        SupPid,
+                        {cowboy_protocol,
+                         {cowboy_protocol, start_link, [Ref, Sock, Transport, CowboyOpts]},
+                         intrinsic, ?MAX_WAIT, worker, [cowboy_protocol]}),
+    {ok, SupPid, ReaderPid}.
+
+start_keepalive_link() ->
+    supervisor2:start_link(?MODULE, []).
+
+%%----------------------------------------------------------------------------
+
+init([]) ->
+    {ok, {{one_for_all, 0, 1}, []}}.
diff --git a/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_handler.erl b/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_handler.erl
@@ -22,103 +22,131 @@
 -export([websocket_info/3]).
 -export([websocket_terminate/3]).
 
+-export([conserve_resources/3]).
+
 -include_lib("amqp_client/include/amqp_client.hrl").
 
 -record(state, {
     conn_name,
+    keepalive,
+    keepalive_sup,
     parse_state,
     proc_state
 }).
 
 init(_, _, _) ->
     {upgrade, protocol, cowboy_websocket}.
 
-websocket_init(_, Req, _) ->
-%    io:format(user, "~p~n", [Req]),
-
+websocket_init(_, Req, Opts) ->
     process_flag(trap_exit, true),
+    {_, KeepaliveSup} = lists:keyfind(keepalive_sup, 1, Opts),
     Sock = cowboy_req:get(socket, Req),
     case rabbit_net:connection_string(Sock, inbound) of
         {ok, ConnStr} ->
             rabbit_log:log(connection, info, "accepting WEB-MQTT connection ~p (~s)~n", [self(), ConnStr]),
-%            rabbit_alarm:register(
-%              self(), {?MODULE, conserve_resources, []}),
+            rabbit_alarm:register(
+                self(), {?MODULE, conserve_resources, []}),
             ProcessorState = rabbit_mqtt_processor:initial_state(Sock,
                 rabbit_mqtt_reader:ssl_login_name(Sock),
                 fun send_reply/2),
             {SecWsProtocol, Req1} = cowboy_req:header(<<"sec-websocket-protocol">>, Req),
             Req2 = cowboy_req:set_resp_header(<<"sec-websocket-protocol">>, SecWsProtocol, Req1),
-            %% @todo compact?
-            %% control_throttle
-            {ok, Req2, #state{
-                %% keepalive/keepalive_sup
-                conn_name   = ConnStr,
-                parse_state = rabbit_mqtt_frame:initial_state(),
-                proc_state  = ProcessorState
-            }};
+            Req3 = cowboy_req:compact(Req2),
+            %% @todo control_throttle
+            {ok, Req3, #state{
+                conn_name     = ConnStr,
+                keepalive     = {none, none},
+                keepalive_sup = KeepaliveSup,
+                parse_state   = rabbit_mqtt_frame:initial_state(),
+                proc_state    = ProcessorState
+            }, hibernate};
         _ ->
             {shutdown, Req}
     end.
 
-%% @todo hibernate everywhere?
 websocket_handle({binary, Data}, Req, State) ->
     handle_data(Data, Req, State);
 websocket_handle(Frame, Req, State) ->
-    rabbit_log:info("rabbit_web_mqtt: unexpected Websocket frame ~p~n",
+    rabbit_log:log(connection, info, "WEB-MQTT: unexpected Websocket frame ~p~n",
                     [Frame]),
-    {ok, Req, State}.
+    {ok, Req, State, hibernate}.
 
 websocket_info({#'basic.deliver'{}, #amqp_msg{}, _DeliveryCtx} = Delivery,
             Req, State = #state{ proc_state = ProcState0 }) ->
     case rabbit_mqtt_processor:amqp_callback(Delivery, ProcState0) of
         {ok, ProcState} ->
-            {ok, Req, State #state { proc_state = ProcState }};
+            {ok, Req, State #state { proc_state = ProcState }, hibernate};
         {error, _, _} ->
             {shutdown, Req, State}
     end;
 websocket_info(#'basic.ack'{} = Ack, Req, State = #state{ proc_state = ProcState0 }) ->
     case rabbit_mqtt_processor:amqp_callback(Ack, ProcState0) of
         {ok, ProcState} ->
-            {ok, Req, State #state { proc_state = ProcState }};
+            {ok, Req, State #state { proc_state = ProcState }, hibernate};
         {error, _, _} ->
             {shutdown, Req, State}
     end;
 websocket_info(#'basic.consume_ok'{}, Req, State) ->
-    {ok, Req, State};
+    {ok, Req, State, hibernate};
 websocket_info(#'basic.cancel'{}, Req, State) ->
     {shutdown, Req, State};
 websocket_info({reply, Data}, Req, State) ->
-    {reply, {binary, Data}, Req, State};
-websocket_info({'EXIT', _, _}, State) ->
+    {reply, {binary, Data}, Req, State, hibernate};
+websocket_info({'EXIT', _, _}, Req, State) ->
     {shutdown, Req, State};
 websocket_info({'$gen_cast', duplicate_id}, Req, State = #state{ proc_state = ProcState,
                                                                  conn_name = ConnName }) ->
-    rabbit_log:warning("MQTT disconnecting duplicate client id ~p (~p)~n",
+    rabbit_log:log(connection, warning, "WEB-MQTT disconnecting duplicate client id ~p (~p)~n",
                  [rabbit_mqtt_processor:info(client_id, ProcState), ConnName]),
     {shutdown, Req, State};
+websocket_info({start_keepalives, Keepalive}, Req,
+               State = #state{ keepalive_sup = KeepaliveSup }) ->
+    Sock = cowboy_req:get(socket, Req),
+    %% Only the client has the responsibility for sending keepalives
+    SendFun = fun() -> ok end,
+    Parent = self(),
+    ReceiveFun = fun() -> Parent ! keepalive_timeout end,
+    Heartbeater = rabbit_heartbeat:start(
+                    KeepaliveSup, Sock, 0, SendFun, Keepalive, ReceiveFun),
+    {ok, Req, State #state { keepalive = Heartbeater }};
+
+websocket_info(keepalive_timeout, Req, State = #state {conn_name = ConnStr,
+                                                       proc_state = PState}) ->
+    rabbit_log:log(connection, error, "closing WEB-MQTT connection ~p (keepalive timeout)~n", [ConnStr]),
+    rabbit_mqtt_processor:send_will(PState),
+    {shutdown, Req, State};
+
+%% @todo conserve_resources
+%% @todo bump_credit
+
 websocket_info(Msg, Req, State) ->
-    rabbit_log:info("rabbit_web_mqtt: unexpected message ~p~n",
+    rabbit_log:log(connection, info, "WEB-MQTT: unexpected message ~p~n",
                     [Msg]),
-    {ok, Req, State}.
-
+    {ok, Req, State, hibernate}.
 
-websocket_terminate(_, _, #state{proc_state=ProcState}) ->
+websocket_terminate(_, _, #state{ proc_state = ProcState,
+                                  conn_name  = ConnName }) ->
+    rabbit_log:log(connection, info, "closing WEB-MQTT connection ~p (~s)~n", [self(), ConnName]),
     rabbit_mqtt_processor:close_connection(ProcState),
     ok;
 websocket_terminate(_, _, _) ->
     ok.
 
+conserve_resources(Pid, _, Conserve) ->
+    Pid ! {conserve_resources, Conserve},
+    ok.
+
 %% Internal.
 
 handle_data(<<>>, Req, State) ->
-    {ok, Req, State};
+    {ok, Req, State, hibernate};
 handle_data(Data, Req, State = #state{ parse_state = ParseState,
                                        proc_state  = ProcState,
                                        conn_name   = ConnStr }) ->
     case rabbit_mqtt_frame:parse(Data, ParseState) of
         {more, ParseState1} ->
             %% @todo control_throttle
-            {ok, Req, State #state{ parse_state = ParseState1 }};
+            {ok, Req, State #state{ parse_state = ParseState1 }, hibernate};
         {ok, Frame, Rest} ->
             case rabbit_mqtt_processor:process_frame(Frame, ProcState) of
                 {ok, ProcState1} ->
@@ -145,5 +173,4 @@ handle_data(Data, Req, State = #state{ parse_state = ParseState,
     end.
 
 send_reply(Frame, _) ->
-    rabbit_log:info("MQTT sending frame ~p ~n", [Frame]),
     self() ! {reply, rabbit_mqtt_frame:serialise(Frame)}.
diff --git a/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_middleware.erl b/deps/rabbitmq_web_mqtt/src/rabbit_web_mqtt_middleware.erl
@@ -0,0 +1,30 @@
+%% The contents of this file are subject to the Mozilla Public License
+%% Version 1.1 (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.mozilla.org/MPL/
+%%
+%% Software distributed under the License is distributed on an "AS IS"
+%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
+%% the License for the specific language governing rights and
+%% limitations under the License.
+%%
+%% The Original Code is RabbitMQ.
+%%
+%% The Initial Developer of the Original Code is GoPivotal, Inc.
+%% Copyright (c) 2016 GoPivotal, Inc.  All rights reserved.
+%%
+
+-module(rabbit_web_mqtt_middleware).
+-behavior(cowboy_middleware).
+
+-export([execute/2]).
+
+execute(Req, Env) ->
+    {keepalive_sup, KeepaliveSup} = lists:keyfind(keepalive_sup, 1, Env),
+    case lists:keyfind(handler_opts, 1, Env) of
+        {_, Opts} when is_list(Opts) ->
+            {ok, Req, lists:keyreplace(handler_opts, 1, Env,
+                {handler_opts, [{keepalive_sup, KeepaliveSup}|Opts]})};
+        _ ->
+            {ok, Req, Env}
+    end.
