Process that monitors the status of the poolboy worker pool used for coverage queries.

pb_status_monitor.erl

%% -------------------------------------------------------------------
%%
%% pb_status_monitor: PoolBoy pool monitor
%%
%% Copyright (c) 2013 Basho Technologies, Inc. All Rights Reserved.
%%
%% This file is provided to you 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(pb_status_monitor).
-compile(export_all).
-behaviour(gen_server).

-export([start/1, start/2, stop/0, is_enabled/0, enable/0, disable/0]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
         terminate/2, code_change/3]).

-record(state, {file, interval}).

-define(DEFAULT_INTERVAL, 10).


%% @doc Check to see if server is enabled.
-spec is_enabled() -> boolean().
is_enabled() ->
    case application:get_env(pb_status_monitor, enabled) of
        {ok, true} ->
            true;
        _ ->
            false
    end.

%% @doc Enable RiakGraphite integration server.
-spec enable() -> ok | {error, term()}.
enable() ->
    case application:get_env(pb_status_monitor, enabled) of
        {ok, error} ->
            {error, "pb_status_monitor encountered error during startup and can not be enabled"};
        {ok, true} ->
            ok;
        {ok, false} ->
            application:set_env(pb_status_monitor, enabled, true),
            ok;
        undefined ->
            {error, "pb_status_monitor has not been started"}
    end.

%% @doc Disable RiakGraphite integration server.
-spec disable() -> ok | {error, term()}.
disable() ->
    case application:get_env(pb_status_monitor, enabled) of
        {ok, error} ->
            {error, "pb_status_monitor encountered error during startup and can not be disabled"};
        {ok, false} ->
            ok;
        {ok, true} ->
            application:set_env(pb_status_monitor, enabled, false),
            ok;
        undefined ->
            {error, "pb_status_monitor has not been started"}
    end.

%% @doc Start PoolBoy Status Monitor server.
-spec start(string()) -> {ok, pid()} | ignore | {error, term()}.
start(File) ->
    start(File, ?DEFAULT_INTERVAL).

%% @doc Start PoolBoy Status Monitor server.
-spec start(string(), pos_integer()) -> {ok, pid()} | ignore | {error, term()}.
start(File, Interval) when is_list(File) andalso is_integer(Interval) andalso Interval > 0 ->
    gen_server:start({local, ?MODULE}, ?MODULE, [File, Interval], []).

%% @doc Stop PoolBoy Status Monitor server.
-spec stop() -> ok.
stop() ->
    gen_server:cast(?MODULE, stop).

%% --------------------------------------------------------------------
%% Function: init/1
%% Description: Initiates the server
%% Returns: {ok, State}          |
%%          {ok, State, Timeout} |
%%          ignore               |
%%          {stop, Reason}
%% --------------------------------------------------------------------    
init([File, Interval]) ->
    case file:open(File, [write]) of 
        {ok, IoDev} ->
            application:set_env(pb_status_monitor, enabled, true),
            erlang:send_after(1000 * Interval, self(), gather),
            {ok, #state{interval = Interval, file = IoDev}};
        {error, Reason} ->
            application:set_env(pb_status_monitor, enabled, error),
            {stop, {error, Reason}}
    end.

%% --------------------------------------------------------------------
%% Function: handle_call/3
%% Description: Handling call messages
%% Returns: {reply, Reply, State}          |
%%          {reply, Reply, State, Timeout} |
%%          {noreply, State}               |
%%          {noreply, State, Timeout}      |
%%          {stop, Reason, Reply, State}   | (terminate/2 is called)
%%          {stop, Reason, State}            (terminate/2 is called)
%% --------------------------------------------------------------------
handle_call(_Msg, _From, State) ->
    {noreply, State}.

%% --------------------------------------------------------------------
%% Function: handle_cast/2
%% Description: Handling cast messages
%% Returns: {noreply, State}          |
%%          {noreply, State, Timeout} |
%%          {stop, Reason, State}            (terminate/2 is called)
%% --------------------------------------------------------------------
handle_cast(stop, State) ->
    {stop, "Process ordered to stop", State};
handle_cast(_Msg, State) ->
    {noreply, State}.

%% --------------------------------------------------------------------
%% Function: handle_info/2
%% Description: Handling all non call/cast messages
%% Returns: {noreply, State}          |
%%          {noreply, State, Timeout} |
%%          {stop, Reason, State}            (terminate/2 is called)
%% --------------------------------------------------------------------
handle_info(gather, #state{interval = Interval, file = File} = State) ->
    case application:get_env(pb_status_monitor, enabled) of
        {ok, true} ->
            VnodePids = [Pid || {_, Pid} <- riak_core_vnode_manager:all_index_pid(riak_kv_vnode)],
            Links = [process_info(Pid, [links]) || Pid <- VnodePids],
            WorkerPoolPids = [WPPid || [{links,[_, WPPid]}] <- Links],
            WorkerPoolLinks = [process_info(Pid, [links]) || Pid <- WorkerPoolPids],
            PoolboyPids = [PoolboyPid || [{links,[_, PoolboyPid]}] <- WorkerPoolLinks],
            List = [poolboy:status(Pid) || Pid <- PoolboyPids],
            Result = dict:to_list(lists:foldl(fun(Item, Acc) -> 
                                                  dict:update_counter(Item, 1, Acc)
                                              end, dict:new(), List)),
            write_to_file(File, Result),
            erlang:send_after(1000 * Interval, self(), gather);
        _ ->
            erlang:send_after(1000 * Interval, self(), gather)
    end,
    {noreply, State};
handle_info(_Info, State) ->
    {noreply, State}.

%% --------------------------------------------------------------------
%% Function: terminate/2
%% Description: Shutdown the server
%% Returns: any (ignored by gen_server)
%% --------------------------------------------------------------------
terminate(_Reason, #state{file = File}) ->
    file:close(File),
    application:set_env(pb_status_monitor, enabled, false),
    ok.

%% --------------------------------------------------------------------
%% Func: code_change/3
%% Purpose: Convert process state when code is changed
%% Returns: {ok, NewState}
%% --------------------------------------------------------------------
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%% Internal functions
write_to_file(File, Result) ->
    LocalTime = erlang:localtime(),
    Str = httpd_util:rfc1123_date(LocalTime),
    Output = io_lib:format("~p: ~p\n\n",[Str, Result]),
    file:write(File, Output).