Perform ring claim dry run, reporting new ring and would-be pending transfers

claim_dryrun.erl

%% Claim dry-run / simulation code
%%
%% If running outside attached Riak shell, must have Riak libs in codepath and
%% proper cookie / long-form node name. Eg:
%% erl -pa RIAK/deps/*/ebin -setcookie riak -name [email protected]
%%
%% Run node/2 passing in cluster claimaint and number of new nodes to test:
%% Add 1 node:  claim_dryrun:node('[email protected]', 1)
%% Add 2 nodes: claim_dryrun:node('[email protected]', 2)
%% ...

-module(claim_dryrun).
-export([node/2]).

node(Claimant, NumToAdd) ->
    {ok, Ring0} = rpc:call(Claimant, riak_core_ring_manager, get_raw_ring, []),
    TargetN = rpc:call(Claimant, app_helper, get_env, [riak_core, target_n_val]),
    application:set_env(riak_core, target_n_val, TargetN),

    case riak_core_ring:pending_changes(Ring0) of
        [] ->
            io:format("Current ring:~n"),
            riak_core_ring:pretty_print(Ring0, [legend]),
            lists:foldl(
              fun(N, Ring) ->
                      Node =
                          list_to_atom(lists:flatten(io_lib:format("[email protected]",[N]))),
                      io:format("~nAdding ~p~n", [Node]),
                      Ring2 = riak_core_ring:add_member(Node, Ring, Node),
                      Ring3 = claim(Claimant, Ring2),
                      Ring3
              end, Ring0, lists:seq(1,NumToAdd));
        _ ->
            io:format("Cannot perform a dryrun on a cluster with pending transfers~n")
    end,
    ok.

claim(Claimant, Ring) ->
    Members = riak_core_ring:claiming_members(Ring),
    Ring2 = lists:foldl(fun(Node, Ring0) ->
                                rpc:call(Claimant,
                                         riak_core_gossip,
                                         claim_until_balanced,
                                         [Ring0, Node])
                        end, Ring, Members),
    Owners1 = riak_core_ring:all_owners(Ring),
    Owners2 = riak_core_ring:all_owners(Ring2),
    Owners3 = lists:zip(Owners1, Owners2),
    Next = [{Idx, PrevOwner, NewOwner}
            || {{Idx, PrevOwner}, {Idx, NewOwner}} <- Owners3,
               PrevOwner /= NewOwner],
    Tally =
        lists:foldl(fun({_, PrevOwner, NewOwner}, Tally) ->
                            dict:update_counter({PrevOwner, NewOwner}, 1, Tally)
                    end, dict:new(), Next),
    riak_core_ring:pretty_print(Ring2, [legend]),
    io:format("Pending: ~p~n", [length(Next)]),
    io:format("Check: ~p~n", [rpc:call(Claimant,
                                       riak_core_ring_util,
                                       check_ring,
                                       [Ring2])]),
    [io:format("~b transfers from ~p to ~p~n", [Count, PrevOwner, NewOwner])
     || {{PrevOwner, NewOwner}, Count} <- dict:to_list(Tally)],
    Ring2.