RabbitMQ的订阅

-module(market_dispatcher_amqp_subscriber).

-behaviour(gen_server).

%% API
-export([start_link/1]).
-export([get_md_statistic/0, reg_websocket/1, unreg_websocket/1]).

%% gen_server callbacks
-export([init/1, 
         handle_call/3, 
         handle_cast/2, 
         handle_info/2,
         terminate/2, 
         code_change/3]).

-define(SERVER, ?MODULE). 
-include_lib("amqp_client/include/amqp_client.hrl").


-record(state, {channel,
                timer,
                clients,
                node_hour_distribution}). %orddict {key:{Node, Time}, value:{ClientCount, MaxInstrumentCount}

%%%===================================================================
%%% API
%%%===================================================================

%%--------------------------------------------------------------------
%% @doc
%% Starts the server
%%
%% @spec start_link() -> {ok, Pid} | ignore | {error, Error}
%% @end
%%--------------------------------------------------------------------
start_link(Params) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [Params], []).

get_md_statistic() ->
    gen_server:call(?SERVER, get_md_statistic).

reg_websocket(Pid) ->
    gen_server:cast(?SERVER, {reg_websocket, Pid}).

unreg_websocket(Pid) ->
    gen_server:cast(?SERVER, {unreg_websocket, Pid}).

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Initializes the server
%%
%% @spec init(Args) -> {ok, State} |
%%                     {ok, State, Timeout} |
%%                     ignore |
%%                     {stop, Reason}
%% @end
%%--------------------------------------------------------------------
init([Params]) ->
    Exchange = config_val(exchange, Params, <<"market_subscriber_statistic">>),
    AmqpParams = #amqp_params_network {
      username       = config_val(amqp_user, Params, <<"guest">>),
      password       = config_val(amqp_pass, Params, <<"guest">>),
      virtual_host   = config_val(amqp_vhost, Params, <<"/">>),
      host           = config_val(amqp_host, Params, "localhost"),
      port           = config_val(amqp_port, Params, 5672)
     },

    Channel = case amqp_channel(AmqpParams) of
                    {ok, Channel0} ->
                        #'exchange.declare_ok'{} = amqp_channel:call(Channel0, #'exchange.declare'{ exchange = Exchange, 
                                                                                                   type = <<"topic">>}),
                        %% Declare a queue
                        #'queue.declare_ok'{queue = Q} = amqp_channel:call(Channel0, #'queue.declare'{queue = <<"md_stat">>}),
                        Binding = #'queue.bind'{queue = Q, exchange = Exchange, routing_key = <<"md_stat">>},
                         #'queue.bind_ok'{} = amqp_channel:call(Channel0, Binding),
                        Sub = #'basic.consume'{queue = Q},
                        % Subscribe the channel and consume the message
                        Consumer = self(),
                        #'basic.consume_ok'{} = amqp_channel:subscribe(Channel0, Sub, Consumer),
                        Channel0;
                    {error, Reason} ->
                        lager:error("Amqp channel start error ~p", [Reason]),
                        undefined
              end,
    {Date, _} = erlang:localtime(),
    EndTime = (calendar:datetime_to_gregorian_seconds({Date, {24, 0, 0}}) - 
               calendar:datetime_to_gregorian_seconds(erlang:localtime())) * 1000,
    {ok, Timer} = timer:send_after(EndTime, clear_charts),

    {ok, #state{channel = Channel,
                timer   = Timer,
                clients = sets:new(), % the pid of websocket
                node_hour_distribution = orddict:new() %orddict {key:{Node, Time}, value:{ClientCount, InstrumentCount}}.
                }}. 

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling call messages
%%
%% @spec handle_call(Request, From, State) ->
%%                                   {reply, Reply, State} |
%%                                   {reply, Reply, State, Timeout} |
%%                                   {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, Reply, State} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_call(get_md_statistic, _From, #state{node_hour_distribution = MDStatistic} 
        = State) ->
    MDStatistic0 = [ [Node, Time, ClientCount, InstrumentCount] 
                || {{Node, Time}, {ClientCount, InstrumentCount}} <- orddict:to_list(MDStatistic)],
    {reply, MDStatistic0, State};


handle_call(_Request, _From, State) ->
    lager:warning("Can't handle request: ~p", [_Request]),
    {reply, {error, invalid_req}, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling cast messages
%%
%% @spec handle_cast(Msg, State) -> {noreply, State} |
%%                                  {noreply, State, Timeout} |
%%                                  {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_cast({reg_websocket, Pid}, #state{clients = Clients} = State) ->
    NewClients = sets:add_element(Pid, Clients),
    {noreply, State#state{clients=NewClients}};

handle_cast({unreg_websocket, Pid}, #state{clients = Clients} = State) ->
    case sets:is_element(Pid, Clients) of
        true  ->
             NewClients = sets:del_element(Pid, Clients),
             {noreply, State#state{clients=NewClients}};
        false ->
             {noreply, State}
    end;

handle_cast(_Msg, State) ->
    lager:warning("Can't handle msg: ~p", [_Msg]),
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling all non call/cast messages
%%
%% @spec handle_info(Info, State) -> {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_info(#'basic.consume_ok'{}, State) ->
    {noreply, State};

handle_info(#'basic.cancel_ok'{}, State) ->
    {noreply, State};

handle_info({#'basic.deliver'{delivery_tag = Tag}, 
    {_, _, Message} = _Content}, #state{} = State) ->
    #state{channel = Channel, clients = Clients,
           node_hour_distribution  = NodeClientInstrumentCount} = State,
    [Node, Time, ClientCount0, InstrumentCount0] = jsx:json_to_term(Message),
    ClientCount =   case ClientCount0 =< 0 of
                        true  -> 0;
                        false -> ClientCount0
                    end, 
    InstrumentCount = case InstrumentCount0 =< 0 of
                        true  -> 0;
                        false -> InstrumentCount0
                      end,
    {NodeClientInstrumentCount0, Msg} = 
    case orddict:find({Node, Time}, NodeClientInstrumentCount) of
        {ok, {_OldClientCount, OldInstrumentCount}} ->
          {NCount, Message1} = case InstrumentCount > OldInstrumentCount of
                            true ->
                                NCount0 = orddict:store({Node, Time}, {ClientCount, InstrumentCount}, 
                                                          NodeClientInstrumentCount),
                                Message0 = {market_dispatcher, [Node, Time, ClientCount, InstrumentCount]},
                                {NCount0, Message0};
                            false ->
                                NCount0 = orddict:store({Node, Time}, {ClientCount, OldInstrumentCount}, 
                                                          NodeClientInstrumentCount),
                                Message0 = {market_dispatcher, [Node, Time, ClientCount, OldInstrumentCount]},
                                {NCount0, Message0}
                          end,
          {NCount,Message1};
        error ->
              NCount  = orddict:store({Node, Time}, {ClientCount, InstrumentCount}, NodeClientInstrumentCount),
              Message0 = {market_dispatcher, [Node, Time, ClientCount, InstrumentCount]},
              {NCount, Message0}
    end,     
    sets:fold(fun(Client, ok) -> 
       Client ! Msg, ok 
     end, ok, Clients),
    amqp_channel:cast(Channel, #'basic.ack'{delivery_tag = Tag}),
    {noreply, State#state{node_hour_distribution = NodeClientInstrumentCount0}};

handle_info(clear_charts, #state{} = State) ->
    #state{timer =Timer} = State,
    {ok, cancel} = timer:cancel(Timer),
    {Date, _} = erlang:localtime(),
    EndTime = ( calendar:datetime_to_gregorian_seconds({Date, {24, 0, 0}}) - 
                calendar:datetime_to_gregorian_seconds(erlang:localtime())) * 1000,
    {ok, Timer2}=timer:send_after(EndTime, clear_charts),

    {noreply, State#state{timer = Timer2,
                          node_hour_distribution      = orddict:new()}};

handle_info(_Info, State) ->
    lager:warning("Can't handle info: ~p", [_Info]),
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% This function is called by a gen_server when it is about to
%% terminate. It should be the opposite of Module:init/1 and do any
%% necessary cleaning up. When it returns, the gen_server terminates
%% with Reason. The return value is ignored.
%%
%% @spec terminate(Reason, State) -> void()
%% @end
%%--------------------------------------------------------------------
terminate(_Reason, _State) ->
    ok.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Convert process state when code is changed
%%
%% @spec code_change(OldVsn, State, Extra) -> {ok, NewState}
%% @end
%%--------------------------------------------------------------------
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%%===================================================================
%%% Internal functions
%%%===================================================================
amqp_channel(AmqpParams) ->
    case maybe_new_pid({AmqpParams, connection},
                       fun() -> amqp_connection:start(AmqpParams) end) of
        {ok, Client} ->
            maybe_new_pid({AmqpParams, channel},
                          fun() -> amqp_connection:open_channel(Client) end);
        Error ->
            Error
    end.

maybe_new_pid(Group, StartFun) ->
    case pg2:get_closest_pid(Group) of
        {error, {no_such_group, _}} ->
            pg2:create(Group),
            maybe_new_pid(Group, StartFun);
        {error, {no_process, _}} ->
            case StartFun() of
                {ok, Pid} ->
                    pg2:join(Group, Pid),
                    {ok, Pid};
                Error ->
                    Error
            end;
        Pid ->
            {ok, Pid}
    end.


config_val(C, Params, Default) ->
  case lists:keyfind(C, 1, Params) of
    {C, V} -> V;
    _ -> Default
  end.