First version of distributed work basho_bench - it needs more work, this was primarily to get something work, and (fixed) remote code loading because traditional erl_prim_loader doesn't jive super well with escript.

Author: sargun

src/basho_bench.erl

@@ -22,7 +22,6 @@
 -module(basho_bench).
 
 -export([main/1, md5/1]).
-
 -include("basho_bench.hrl").
 
 %% ====================================================================
@@ -53,6 +52,8 @@ main(Args) ->
         {error, {already_loaded, basho_bench}} -> ok
     end,
     register(basho_bench, self()),
+    %% TODO: Move into a proper supervision tree, janky for now
+    {ok, _Pid} = basho_bench_config:start_link(),
     basho_bench_config:set(test_id, BenchName),
 
     application:load(lager),
@@ -93,15 +94,16 @@ main(Args) ->
     %% Copy the config into the test dir for posterity
     [ begin {ok, _} = file:copy(Config, filename:join(TestDir, filename:basename(Config))) end
       || Config <- Configs ],
-
+    case basho_bench_config:get(distribute_work, false) of 
+        true -> setup_distributed_work();
+        false -> ok
+    end,
     %% Set our CWD to the test dir
     ok = file:set_cwd(TestDir),
-
     log_dimensions(),
 
     %% Run pre_hook for user code preconditions
     run_pre_hook(),
-
     %% Spin up the application
     ok = basho_bench_app:start(),
 
@@ -276,6 +278,84 @@ run_hook({Module, Function}) ->
 run_hook(no_op) ->
     no_op.
 
+get_addr_args() ->
+    {ok, IfAddrs} = inet:getifaddrs(),
+    FlattAttrib = lists:flatten([IfAttrib || {_Ifname, IfAttrib} <- IfAddrs]),
+    Addrs = proplists:get_all_values(addr, FlattAttrib),
+    StrAddrs = lists:map(fun(X) -> inet:ntoa(X) end, Addrs),
+    string:join(StrAddrs, " ").
+setup_distributed_work() ->
+    io:format("Using cookie: ~p~n", [erlang:get_cookie()]),
+    case node() of 
+        'nonode@nohost' -> 
+            ?STD_ERR("Basho bench not started in distributed mode, and distribute_work = true~n", []),
+            halt(1);
+        _ -> ok
+    end,
+    {ok, _Pid} = erl_boot_server:start([]),
+    %% Allow anyone to boot from me...I might want to lock this down this down at some point
+    erl_boot_server:add_subnet({0,0,0,0}, {0,0,0,0}),
+    %% This is cheating, horribly, but it's the only simple way to bypass net_adm:host_file()
+    gen_server:start({global, pool_master}, pool, [], []),
+    RemoteSpec = basho_bench_config:get(remote_nodes, []),
+    Cookie = lists:flatten(erlang:atom_to_list(erlang:get_cookie())),
+    Args = "-setcookie " ++ Cookie ++ " -loader inet -hosts " ++ get_addr_args(),
+    [io:format("Starting slave {~p, ~p, ~p}~n", [Host, Name, Args]) || {Host, Name} <- RemoteSpec],
+    Slaves = [ slave:start_link(Host, Name, Args) || {Host, Name} <- RemoteSpec],
+    SlaveNames = [SlaveName || {ok, SlaveName} <- Slaves],
+    [pool:attach(SlaveName) || SlaveName <- SlaveNames],
+    CodePaths = code:get_path(),
+    rpc:multicall(SlaveNames, code, set_path, [CodePaths]),
+    Apps = [lager, basho_bench, getopt, bear, folsom, ibrowse, riakc, riak_pb, mochiweb, protobuffs, velvet, goldrush],
+    [distribute_app(App, SlaveNames) || App <- Apps].
+  %  io:format("Running escript: ~p~n", [escript:script_name()]),
+    %% Assumes I'm running basho_bench from an escript, otherwise the behaviour is undefined
+%    EscriptFileName = escript:script_name(),
+%    {ok, Sections} = escript:extract(EscriptFileName, [compile_source]),
+%    io:format("Sections: ~p~n", [Sections]),
+
+%    Libdir = code:lib_dir(basho_bench),
+ %   LibdirLen = string:len(Libdir), 
+  %  BashoBenchLibs = lists:filter(fun(X) -> string:substr(X, 1, LibdirLen) == Libdir end, code:get_path()), 
+   % io:format("Basho Bench Library Directoryies: ~p~n", [BashoBenchLibs]),
+%    io:format("Basho_bench lib dir: ~p~n", [Libdir]),
+%    io:format("Libdir: ~p~n", [erl_prim_loader:list_dir(Libdir)]),
+%    io:format("Code Path: ~p~n", [code:get_path()]),
+
+%    io:format("Files? ~p~n", [erl_prim_loader:list_dir("/Users/sdhillon/repos/basho_bench/basho_bench/basho_bench/ebin")]),
+%    io:format("Remote nodes: ~p~n", [nodes()]),
+%    RemoteFiles = rpc:call('bb25@3c075477e55e-2.local', erl_prim_loader, list_dir, ["/Users/sdhillon/repos/basho_bench/basho_bench/basho_bench/ebin"]),  
+%    io:format("Remote files: ~p~n", [RemoteFiles]),
+
+
+deploy_module(Module, Nodes) ->
+    {Module, Binary, Filename} = code:get_object_code(Module),
+    rpc:multicall(Nodes, code, load_binary, [Module, Filename, Binary]).
+
+distribute_app(App, Nodes) ->
+    % :(. This is super hackish, it depends on a bunch of assumptions
+    % But, unfortunately there are negative interactions with escript and slave nodes
+    CodeExtension = code:objfile_extension(),
+    LibDir = code:lib_dir(App),
+    % Get what paths are in the code path that start with LibDir
+    LibDirLen = string:len(LibDir),
+    EbinsDir = lists:filter(fun(CodePathDir) -> string:substr(CodePathDir, 1, LibDirLen) ==  LibDir end, code:get_path()),
+    StripEndFun = fun(Path) ->
+        PathLen = string:len(Path),
+        case string:substr(Path, PathLen - string:len(CodeExtension) + 1, string:len(Path)) of 
+            CodeExtension ->
+                {true, string:substr(Path, 1, PathLen - string:len(CodeExtension))};
+            _ -> false
+        end
+    end, 
+    EbinDirDistributeFun = fun(EbinDir) ->
+        {ok, Beams} = erl_prim_loader:list_dir(EbinDir),
+        Modules = lists:filtermap(StripEndFun, Beams),
+        ModulesLoaded = lists:map(fun(X) -> code:load_abs(filename:join(EbinDir, X)) end, Modules),
+        lists:foreach(fun({module, Module}) -> deploy_module(Module, Nodes) end, ModulesLoaded)
+    end,
+    lists:foreach(EbinDirDistributeFun, EbinsDir),    
+    ok.
 %% just a utility, should be in basho_bench_utils.erl
 %% but 's' is for multiple utilities, and so far this
 %% is the only one.

src/basho_bench_config.erl

@@ -20,27 +20,62 @@
 %%
 %% -------------------------------------------------------------------
 -module(basho_bench_config).
+-behaviour(gen_server).
+
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+-compile(export_all).
+-endif.
 
 -export([load/1,
          normalize_ips/2,
          set/2,
          get/1, get/2]).
 
+-export([start_link/0]).
+
+% Gen server callbacks
+-export([code_change/3, init/1, terminate/2, handle_call/3]).
+
 -include("basho_bench.hrl").
 
+-record(basho_bench_config_state, {config_keys}).
+
+-type state() :: #basho_bench_config_state{}.
 %% ===================================================================
 %% Public API
 %% ===================================================================
 
+ensure_started() -> 
+    start_link().
+
+start_link() ->
+    gen_server:start_link({global, ?MODULE}, ?MODULE, [], []).
+
+
 load(Files) ->
-    TermsList =
-        [ case file:consult(File) of
-              {ok, Terms} ->
-                  Terms;
-              {error, Reason} ->
-                  ?FAIL_MSG("Failed to parse config file ~s: ~p\n", [File, Reason])
-          end || File <- Files ],
-    load_config(lists:append(TermsList)).
+    ensure_started(),
+    gen_server:call({global, ?MODULE}, {load_files, Files}). 
+    
+set(Key, Value) ->
+    gen_server:call({global, ?MODULE}, {set, Key, Value}).
+
+get(Key) ->
+    case gen_server:call({global, ?MODULE}, {get, Key}) of
+        {ok, Value} ->
+            Value;
+        error ->
+            erlang:error("Missing configuration key", [Key])
+    end.
+
+get(Key, Default) ->
+    case gen_server:call({global, ?MODULE}, {get, Key}) of
+        {ok, Value} ->
+            Value;
+        error ->
+            Default
+    end.
 
 %% @doc Normalize the list of IPs and Ports.
 %%
@@ -58,42 +93,177 @@ normalize_ips(IPs, DefultPort) ->
         end,
     lists:foldl(F, [], IPs).
 
-set(Key, Value) ->
-    ok = application:set_env(basho_bench, Key, Value).
 
-get(Key) ->
-    case application:get_env(basho_bench, Key) of
-        {ok, Value} ->
-            Value;
-        undefined ->
-            erlang:error("Missing configuration key", [Key])
-    end.
 
-get(Key, Default) ->
-    case application:get_env(basho_bench, Key) of
-        {ok, Value} ->
-            Value;
-        _ ->
-            Default
-    end.
 
 
 %% ===================================================================
 %% Internal functions
 %% ===================================================================
 
-load_config([]) ->
-    ok;
-load_config([{Key, Value} | Rest]) ->
-    ?MODULE:set(Key, Value),
-    load_config(Rest);
-load_config([ Other | Rest]) ->
-    ?WARN("Ignoring non-tuple config value: ~p\n", [Other]),
-    load_config(Rest).
 
 normalize_ip_entry({IP, Ports}, Normalized, _) when is_list(Ports) ->
     [{IP, Port} || Port <- Ports] ++ Normalized;
 normalize_ip_entry({IP, Port}, Normalized, _) ->
     [{IP, Port}|Normalized];
 normalize_ip_entry(IP, Normalized, DefaultPort) ->
     [{IP, DefaultPort}|Normalized].
+
+
+%% ===
+%% Gen_server Functions
+%% ===
+
+-spec init(term()) -> {ok, state()}.  
+init(_Args) ->
+    State = #basho_bench_config_state{config_keys = base_config()},
+    {ok, State}.
+
+-spec code_change(term(), state(), term()) -> {ok, state()}.
+code_change(_OldVsn, State, _Extra) ->
+    {ok, State}.                                
+
+-spec terminate(term(), state()) -> 'ok'.
+terminate(_Reason, _State) ->
+    ok.
+
+handle_call({load_files, FileNames}, _From, State) ->
+  TermsList = get_keys_from_files(FileNames),
+  ConfigKeys2 = load_termlist(TermsList, State#basho_bench_config_state.config_keys),
+  State2 = State#basho_bench_config_state{config_keys = ConfigKeys2},
+  {reply, ok, State2};
+handle_call({set, app_run_mode, Value}, _From, State) ->
+  application:set_env(basho_bench_app, app_run_mode, Value),
+  {reply, ok, State};
+handle_call({get, app_run_mode}, _From, State) ->
+  case application:get_env(basho_bench_app, app_run_mode) of
+    {ok, Value} ->
+      {reply, {ok, Value}, State};
+    undefined ->
+      {reply, error, State}
+  end;
+handle_call({set, Key, Value}, _From, State) ->
+  NewKVs = orddict:store(Key, Value, State#basho_bench_config_state.config_keys),
+  State2 = State#basho_bench_config_state{config_keys = NewKVs},
+  {reply, ok, State2};
+handle_call({get, Key}, _From, State) ->
+  Value = orddict:find(Key, State#basho_bench_config_state.config_keys),
+  {reply, Value, State}.
+get_keys_from_files(Files) ->
+    [ case file:consult(File) of
+          {ok, Terms} ->
+              Terms;
+          {error, Reason} ->
+              ?FAIL_MSG("Failed to parse config file ~s: ~p\n", [File, Reason]),
+              throw(invalid_config),
+              notokay
+      end || File <- Files ].
+
+
+load_termlist(TermList, ExistingConfig) ->
+    NewKVs = lists:flatten(TermList),
+    FoldFun = fun({Key, Value}, Accum) ->
+        orddict:store(Key, Value, Accum)
+    end,
+    lists:foldl(FoldFun, ExistingConfig, NewKVs).
+
+base_config() ->
+  orddict:from_list([
+        %% Run mode: How should basho_bench started as a separate node, or part of an 
+        %% other node. The default is standalone, other option is included.
+   
+         %%
+         %% Mode of load generation:
+         %% max - Generate as many requests as possible per worker
+         %% {rate, Rate} - Exp. distributed Mean reqs/sec
+         %%
+         {mode, {rate, 5}},
+
+         %%
+         %% Default log level
+         %%
+         {log_level, debug},
+
+         %%
+         %% Base test output directory
+         %%
+         {test_dir, "tests"},
+
+         %%
+         %% Test duration (minutes)
+         %%
+         {duration, 5},
+
+         %%
+         %% Number of concurrent workers
+         %%
+         {concurrent, 3},
+
+         %%
+         %% Driver module for the current test
+         %%
+         {driver, basho_bench_driver_http_raw},
+
+         %%
+         %% Operations (and associated mix). Note that
+         %% the driver may not implement every operation.
+         %%
+         {operations, [{get, 4},
+                       {put, 4},
+                       {delete, 1}]},
+
+         %%
+         %% Interval on which to report latencies and status (seconds)
+         %%
+         {report_interval, 10},
+
+         %%
+         %% Key generators
+         %%
+         %% {uniform_int, N} - Choose a uniformly distributed integer between 0 and N
+         %%
+         {key_generator, {uniform_int, 100000}},
+
+         %%
+         %% Value generators
+         %%
+         %% {fixed_bin, N} - Fixed size binary blob of N bytes
+         %%
+         {value_generator, {fixed_bin, 100}}
+        ]).
+-ifdef(TEST).
+load_files_test() ->
+  %% Extracted from bitcask, and null test.
+  KVs = [[{mode,max},
+       {duration,1},
+       {report_interval,1},
+       {concurrent,8},
+       {driver,basho_bench_driver_null},
+       {key_generator,{partitioned_sequential_int,5000000}},
+       {disable_sequential_int_progress_report,true},
+       {value_generator,{fixed_bin,10248}},
+       {operations,[{do_something,7},{an_error,1},{another_error,2}]}],
+      [{mode,max},
+       {duration,10},
+       {concurrent,1},
+       {driver,basho_bench_driver_bitcask},
+       {key_generator,{int_to_bin_bigendian,{uniform_int,5000000}}},
+       {value_generator,{fixed_bin,10000}},
+       {operations,[{get,1},{put,1}]},
+       {code_paths,["../../public/bitcask"]},
+       {bitcask_dir,"/tmp/bitcask.bench"},
+       {bitcask_flags,[o_sync]}]],
+  KVOrdDict = [{bitcask_dir,"/tmp/bitcask.bench"},
+      {bitcask_flags,[o_sync]},
+      {code_paths,["../../public/bitcask"]},
+      {concurrent,1},
+      {disable_sequential_int_progress_report,true},
+      {driver,basho_bench_driver_bitcask},
+      {duration,10},
+      {key_generator,{int_to_bin_bigendian,{uniform_int,5000000}}},
+      {mode,max},
+      {operations,[{get,1},{put,1}]},
+      {report_interval,1},
+      {value_generator,{fixed_bin,10000}}],
+  ?assertEqual(KVOrdDict, load_termlist(KVs, [])).
+-endif.