riak_cs_data_usage.erl

#!/usr/bin/env escript

main(Args) ->

	% - From nodetool ------------------------------------------------------- %

	{_, TargetNode} = process_args(Args, [], undefined),

	case {net_kernel:connect_node(TargetNode), net_adm:ping(TargetNode)} of
        {true, pong} ->
            ok;
        {false,pong} ->
            io:format("Failed to connect to node ~p .\n", [TargetNode]),
            halt(1);
        {_, pang} ->
            io:format("Node ~p not responding to pings.\n", [TargetNode]),
            halt(1)
    end,

    % ---------------------------------------------------------------------- %

    % Get List of {Partition, Owner}, Active nodes and Down nodes.
    Get_Members = fun() ->
    	{ok, Ring} = riak_core_ring_manager:get_my_ring(),
    	All_Owners = riak_core_ring:all_owners(Ring),
    	Active_Members = riak_core_ring:active_members(Ring),
    	Down_Members = riak_core_ring:down_members(Ring),
    	{All_Owners, Active_Members, Down_Members}
    end,

	% Run Get_Members on node that this script is executed on.
    Members = rpc:call(TargetNode, erlang, apply, [Get_Members, []]),

    {All_Owners, Active_Members, Down_Members} = Members,

    % Get list of {Partition, Owner} pairs for only Active nodes.
    Active_Owners = lists:filter(fun(X) -> {_, Owner} = X, lists:member(Owner, Active_Members) end, All_Owners),

    % Print a message if there are any down nodes - this count is not accurate if nodes are down.
    if 
    	length(Down_Members) > 0 ->
    	io:format("The following nodes are down: ~p\nPlease run again when all nodes are UP.\n\n", [Down_Members]);
    	true -> ok
    end,

    Total = lists:foldl(fun(F,A) -> process_owner(F) + A end, 0, Active_Owners),

    io:format("Total Bytes: ~p\n", [Total]).

% For a given Owner, open Bitcask directory and get Bitcask status, 
% getting number of bytes used for each file in the Status.
process_owner(Owner) ->
	{Partition, Node} = Owner,
	Get_Bitcask_Status = fun(Partition2) ->
		{ok, Config} = application:get_env(riak_kv, multi_backend),
		Params = [P || {Name, _ , P} <- Config, Name == <<"be_blocks">> orelse Name == be_blocks],
		[Data_Dir] = [D || [{N, D}] <- Params, N == data_root],
		Partition_Ref = bitcask:open(Data_Dir ++ "/" ++ integer_to_list(Partition2)),
		Status = bitcask:status(Partition_Ref),
		Status
	end,
	Status = rpc:call(Node, erlang, apply, [Get_Bitcask_Status, [Partition]]),
	{_, Data_Files} = Status,
	lists:foldl(fun(F,A) -> process_datafile(F) + A end, 0, Data_Files).

% Get number of bytes being used from a Bitcask data file.
process_datafile(DataFile) ->
	{_,_,Dead_Bytes,Total_Bytes} = DataFile,
	Bytes_Used = Total_Bytes - Dead_Bytes,
	Bytes_Used.

% --------------------------------------------------------------- %
% - Functions from nodetool ------------------------------------- %
% --------------------------------------------------------------- %

process_args([], Acc, TargetNode) ->
    {lists:reverse(Acc), TargetNode};
process_args(["-setcookie", Cookie | Rest], Acc, TargetNode) ->
    erlang:set_cookie(node(), list_to_atom(Cookie)),
    process_args(Rest, Acc, TargetNode);
process_args(["-name", TargetName | Rest], Acc, _) ->
    ThisNode = append_node_suffix(TargetName, "_maint_"),
    {ok, _} = net_kernel:start([ThisNode, longnames]),
    process_args(Rest, Acc, nodename(TargetName));
process_args(["-sname", TargetName | Rest], Acc, _) ->
    ThisNode = append_node_suffix(TargetName, "_maint_"),
    {ok, _} = net_kernel:start([ThisNode, shortnames]),
    process_args(Rest, Acc, nodename(TargetName));
process_args([Arg | Rest], Acc, Opts) ->
    process_args(Rest, [Arg | Acc], Opts).

nodename(Name) ->
    case string:tokens(Name, "@") of
        [_Node, _Host] ->
            list_to_atom(Name);
        [Node] ->
            [_, Host] = string:tokens(atom_to_list(node()), "@"),
            list_to_atom(lists:concat([Node, "@", Host]))
    end.

append_node_suffix(Name, Suffix) ->
    case string:tokens(Name, "@") of
        [Node, Host] ->
            list_to_atom(lists:concat([Node, Suffix, os:getpid(), "@", Host]));
        [Node] ->
            list_to_atom(lists:concat([Node, Suffix, os:getpid()]))
    end.