Escript to query various in-VM metrics on a Riak node not currently available from stats.

riak_metrics.escript

-module(riak_metrics).
-compile(export_all).
 
main([NodeName0, Cookie, Length, Command]) ->
 LocalName = '[email protected]',
 NodeName = list_to_atom(NodeName0),
 case net_kernel:start([LocalName]) of
        {ok, _} ->
            erlang:set_cookie(node(), list_to_atom(Cookie)),
            case net_kernel:hidden_connect_node(NodeName) of
                true ->
                    case list_to_atom(Command)  of
                        registered ->
                            RQ = fun(Max) -> 
                                R = fun(MaxLen) -> 
                                    {node(), [ { Name, Len, Pid } || Pid <- processes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len >= MaxLen, {registered_name, Name} <- [process_info(Pid, registered_name)]]} 
                                end,
                            rpc:multicall(erlang,apply,[R,[Max]])
                            end,
                            Registered = rpc:call(NodeName,erlang,apply,[RQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [Registered]),
                            ok;
                        all ->
                            MQ = fun(Max) -> 
                                M = fun(MaxLen) -> 
                                    {node(), [ { Len, Pid } || Pid <- processes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len >= MaxLen]} 
                                end,
                            rpc:multicall(erlang,apply,[M,[Max]])
                            end,
                            All = rpc:call(NodeName,erlang,apply,[MQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [All]),
                            ok;
                        kv_vnodes ->
                            VQ = fun(Max) -> 
                                V = fun(MaxLen) -> 
                                    {node(), [{Index, Len, Pid} || {riak_kv_vnode, Index, Pid} <- riak_core_vnode_manager:all_vnodes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len >= MaxLen]}
                                end,
                            rpc:multicall(erlang,apply,[V,[Max]])
                            end,
                            Vnodes = rpc:call(NodeName,erlang,apply,[VQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [Vnodes]),
                            ok;
                        poolboy ->
                            VQ = fun(_Max) -> 
                                V = fun() -> 
                                    {node(), [{Index, poolboy:status(PoolPid), PoolPid} || {riak_kv_vnode, Index, Pid} <- riak_core_vnode_manager:all_vnodes(), {links,[_,WorkPoolPid]} <- [process_info(Pid, links)], {links,[_,PoolPid]} <- [process_info(WorkPoolPid, links)]]}
                                end,
                            rpc:multicall(erlang,apply,[V,[]])
                            end,
                            Vnodes = rpc:call(NodeName,erlang,apply,[VQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [Vnodes]),
                            ok;
                        % TO-DO: Function to dump message queue of vnode to file on local node 
                        % dump_kv_vnodes ->
                        %     VQ = fun(Max) -> 
                        %         V = fun(MaxLen) -> 
                        %             {node(), [{Index, Len, Pid} || {riak_kv_vnode, Index, Pid} <- riak_core_vnode_manager:all_vnodes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len >= MaxLen]}
                        %         end,
                        %     rpc:multicall(erlang,apply,[V,[Max]])
                        %     end,
                        %     Vnodes = rpc:call(NodeName,erlang,apply,[VQ,[list_to_integer(Length)]]),
                        %     io:format("~p~n", [Vnodes]),
                        %     ok;
                        hist_kv_vnodes_get ->
                            HQ = fun() -> 
                                H = fun() -> 
                                    {[calendar:local_time(),node()],[{Idx, proplists:get_value(n, Stats),proplists:get_value(percentile, Stats), proplists:get_value(max, Stats)} || {riak_kv_vnode, Idx, _Pid} <- riak_core_vnode_manager:all_vnodes(), (Stats = riak_core_stat_q:calc_stat({{riak_kv,vnode,gets,time,list_to_atom(integer_to_list(Idx))},histogram})) /= undefined]}
                                end,
                            rpc:multicall(erlang,apply,[H,[]])
                            end,
                            Histogram = rpc:call(NodeName,erlang,apply,[HQ,[]]),
                            io:format("~p~n", [Histogram]),
                            ok;
                        hist_kv_vnodes_put ->
                            HQ = fun() -> 
                                H = fun() -> 
                                    {[calendar:local_time(),node()],[{Idx, proplists:get_value(n, Stats),proplists:get_value(percentile, Stats), proplists:get_value(max, Stats)} || {riak_kv_vnode, Idx, _Pid} <- riak_core_vnode_manager:all_vnodes(), (Stats = riak_core_stat_q:calc_stat({{riak_kv,vnode,puts,time,list_to_atom(integer_to_list(Idx))},histogram})) /= undefined]}
                                end,
                            rpc:multicall(erlang,apply,[H,[]])
                            end,
                            Histogram = rpc:call(NodeName,erlang,apply,[HQ,[]]),
                            io:format("~p~n", [Histogram]),
                            ok;
                        kill_kv_vnodes ->
                            KQ = fun(Max) ->
                                K = fun(MaxLen) ->
                                    {node(), [{Index, Pid, exit(Pid, kill)} || {riak_kv_vnode, Index, Pid} <- riak_core_vnode_manager:all_vnodes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len > MaxLen]}
                                end,
                            rpc:multicall(erlang,apply,[K,[Max]])
                            end,
                            Killed = rpc:call(NodeName,erlang,apply,[KQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [Killed]),
                            ok;
                        kill_all ->
                            KQ = fun(Max) ->
                                K = fun(MaxLen) ->
                                    {node(), [{Pid, exit(Pid, kill)} || Pid <- processes(), {message_queue_len,Len} <- [process_info(Pid, message_queue_len)], Len > MaxLen]}
                                end,
                            rpc:multicall(erlang,apply,[K,[Max]])
                            end,
                            Killed = rpc:call(NodeName,erlang,apply,[KQ,[list_to_integer(Length)]]),
                            io:format("~p~n", [Killed]),
                            ok
                    end;
                false ->
                    io:format("Could not connect to ~s with cookie ~s", [NodeName, Cookie]);
                _ ->
                    io:format("net_kernel:connect/1 reports ~s is not alive", [LocalName])
            end; 
        {error, Reason} ->
            io:format("Could not connect node: ~w~n", [Reason])
    end;
main(_) ->
  io:format("Usage: riak escript riak_metrics NODENAME COOKIE QUEUE_LENGTH [registered | all | kv_vnodes | hist_kv_vnodes_get | hist_kv_vnodes_get | kill_kv_vnodes | kill_all]~n").