diff --git a/intercepts/intercept.erl b/intercepts/intercept.erl index 1b9a0aaae..b907ed072 100644 --- a/intercepts/intercept.erl +++ b/intercepts/intercept.erl @@ -1,8 +1,8 @@ -module(intercept). %% Export explicit API but also send compile directive to export all -%% becuase some of these private functions are useful in their own +%% because some of these private functions are useful in their own %% right. --export([add/3]). +-export([add/3, add/4]). -compile(export_all). -type abstract_code() :: term(). @@ -35,7 +35,7 @@ %% %% E.g. `[{{update_perform,2}, sleep_update_perform}]' -spec add(module(), module(), mapping()) -> ok. -add(Target, Intercept, Mapping) -> +add(Target, Intercept, Mapping, OutDir) -> Original = ?ORIGINAL(Target), TargetAC = get_abstract_code(Target), @@ -43,16 +43,22 @@ add(Target, Intercept, Mapping) -> Original, TargetAC), OrigAC = make_orig_abstract_code(Target, Original, TargetAC), - ok = compile_and_load(Original, OrigAC), - ok = compile_and_load(Target, ProxyAC). + ok = compile_and_load(Original, OrigAC, OutDir), + ok = compile_and_load(Target, ProxyAC, OutDir). + +add(Target, Intercept, Mapping) -> + add(Target, Intercept, Mapping, undefined). %% @private %% %% @doc Compile the abstract code `AC' and load it into the code server. --spec compile_and_load(module(), abstract_code()) -> ok. -compile_and_load(Module, AC) -> +-spec compile_and_load(module(), abstract_code(), undefined | string()) -> ok. +compile_and_load(Module, AC, OutDir) -> {ok, Module, Bin} = compile:forms(AC,[debug_info]), - {module, Module} = code:load_binary(Module, atom_to_list(Module), Bin), + ModStr = atom_to_list(Module), + _ = is_list(OutDir) andalso + file:write_file(filename:join(OutDir, ModStr ++ ".beam"), Bin), + {module, Module} = code:load_binary(Module, ModStr, Bin), ok. %% @private diff --git a/src/rt_intercept.erl b/src/rt_intercept.erl index f767a76e9..b8e3ae4ad 100644 --- a/src/rt_intercept.erl +++ b/src/rt_intercept.erl @@ -46,16 +46,24 @@ load_code(Node) -> [ok = remote_compile_and_load(Node, F) || F <- intercept_files()], ok. -add(Node, Intercepts) when is_list(Intercepts) -> - [ok = add(Node, I) || I <- Intercepts], +add_and_save(Node, Intercepts) -> + CodePaths = rpc:call(Node, code, get_path, []), + [PatchesDir] = [P || P <- CodePaths, lists:suffix("basho-patches", P)], + add(Node, Intercepts, PatchesDir). + +add(Node, Intercepts) -> + add(Node, Intercepts, undefined). + +add(Node, Intercepts, OutDir) when is_list(Intercepts) -> + [ok = add(Node, I, OutDir) || I <- Intercepts], ok; -add(Node, {Target, Mapping}) -> - add(Node, {Target, ?DEFAULT_INTERCEPT(Target), Mapping}); +add(Node, {Target, Mapping}, OutDir) -> + add(Node, {Target, ?DEFAULT_INTERCEPT(Target), Mapping}, OutDir); -add(Node, {Target, Intercept, Mapping}) -> +add(Node, {Target, Intercept, Mapping}, OutDir) -> NMapping = [transform_anon_fun(M) || M <- Mapping], - ok = rpc:call(Node, intercept, add, [Target, Intercept, NMapping]). + ok = rpc:call(Node, intercept, add, [Target, Intercept, NMapping, OutDir]). %% The following function transforms anonymous function mappings passed %% from an Erlang shell. Anonymous intercept functions from compiled code diff --git a/src/rt_intercept_pt.erl b/src/rt_intercept_pt.erl index 2d145c701..0a58555f3 100644 --- a/src/rt_intercept_pt.erl +++ b/src/rt_intercept_pt.erl @@ -66,7 +66,8 @@ form({clause,L,H,G,B}) -> {clause,L,H,G,forms(B)}; form({match,L,Lhs,Rhs}) -> {match,L,forms(Lhs),forms(Rhs)}; -form({call,L,{remote,_,{atom,_,rt_intercept},{atom,_,add}}=Fun,Args}) -> +form({call,L,{remote,_,{atom,_,rt_intercept},{atom,_,AddFunction}}=Fun,Args}) + when AddFunction == add; AddFunction == add_and_save -> [Node, Intercept] = Args, {call,L,Fun,[Node,intercept(Intercept)]}; form(F) when is_tuple(F) -> diff --git a/tests/ensemble_basic2.erl b/tests/ensemble_basic2.erl index dbcf515a6..dee9e815f 100644 --- a/tests/ensemble_basic2.erl +++ b/tests/ensemble_basic2.erl @@ -20,6 +20,7 @@ -module(ensemble_basic2). -export([confirm/0]). +-compile({parse_transform, rt_intercept_pt}). -include_lib("eunit/include/eunit.hrl"). confirm() -> @@ -33,4 +34,43 @@ confirm() -> lager:info("Killing all ensemble leaders"), ok = ensemble_util:kill_leaders(Node, Ensembles), ensemble_util:wait_until_stable(Node, NVal), + Peers = [PeerId || {PeerId, _PeerPid} <- ensemble_util:peers(Node)], + lager:info("Verifying peers wait for riak_kv_service"), + Delay = rt_config:get(kv_vnode_delay, 5000), + rt_intercept:add_and_save(Node, {riak_kv_vnode, [{{init, 1}, {[Delay], + fun(Args) -> + timer:sleep(Delay), + riak_kv_vnode_orig:init_orig(Args) + end}}]}), + rt:stop_and_wait(Node), + rt:start(Node), + lager:info("Polling peers while riak_kv starts. We should see none"), + UpNoPeersFun = + fun() -> + PL = ensemble_util:peers(Node), + NodePeers = [P || {P, _} <- PL], + NonRootPeers = [P || P <- NodePeers, element(1, P) /= root], + S = rpc:call(Node, riak_core_node_watcher, services, [Node]), + case S of + L when is_list(L) -> + case lists:member(riak_kv, L) of + true -> + true; + false -> + ?assertEqual([], NonRootPeers) + end; + Err -> + ?assertEqual(ok, {peer_get_error, Err}) + end + end, + rt:wait_until(UpNoPeersFun), + lager:info("Perfect. riak_kv is now up and no peers started before that. " + "Now check they come back up"), + SPeers = lists:sort(Peers), + ?assertEqual(ok, rt:wait_until(fun() -> + L = ensemble_util:peers(Node), + L2 = lists:sort([P || {P, _} <- L]), + SPeers == L2 + end)), + lager:info("All expected peers are back. Life is good"), pass. diff --git a/tests/ensemble_util.erl b/tests/ensemble_util.erl index 1c5dfc632..d6f79145b 100644 --- a/tests/ensemble_util.erl +++ b/tests/ensemble_util.erl @@ -79,6 +79,9 @@ ensembles(Node) -> get_leader_pid(Node, Ensemble) -> rpc:call(Node, riak_ensemble_manager, get_leader_pid, [Ensemble]). +peers(Node) -> + rpc:call(Node, riak_ensemble_peer_sup, peers, []). + kill_leader(Node, Ensemble) -> case get_leader_pid(Node, Ensemble) of undefined ->