Add alternative unique PUTs

One which will check that the PUT has replicated, another to replicate bet365's switching buckets and lack of indexes.

Author: martinsumner

src/basho_bench_driver_nhs.erl

@@ -32,6 +32,7 @@
 
 -record(state, {
                 pb_pid,
+                repl_pid,
                 http_host,
                 http_port,
                 recordBucket,
@@ -98,6 +99,7 @@ new(Id) ->
     PBPort  = basho_bench_config:get(pb_port, 8087),
     HTTPIPs = basho_bench_config:get(http_ips, ["127.0.0.1"]),
     HTTPPort =  basho_bench_config:get(http_port, 8098),
+    ReplPBIPs = basho_bench_config:get(replpb_ips, ["127.0.0.1"]),
 
     PBTimeout = basho_bench_config:get(pb_timeout_general, 30*1000),
     HTTPTimeout = basho_bench_config:get(http_timeout_general, 30*1000),
@@ -120,6 +122,12 @@ new(Id) ->
     ?INFO("Using pb target ~p:~p for worker ~p\n", [PBTargetIp,
                                                     PBTargetPort,
                                                     Id]),
+    ReplTargets = basho_bench_config:normalize_ips(ReplPBIPs, PBPort),
+    {ReplTargetIp,
+        ReplTargetPort} = lists:nth((Id rem length(ReplTargets) + 1),
+                                    ReplTargets),
+    ?INFO("Using repl target ~p:~p for worker ~p\n", 
+            [ReplTargetIp, ReplTargetPort, Id]),
 
     {AGMaxKC, AGMinKC, AGKeyOrder} = 
         basho_bench_config:get(alwaysget, {1, 1, key_order}),
@@ -134,8 +142,17 @@ new(Id) ->
     case riakc_pb_socket:start_link(PBTargetIp, PBTargetPort) of
         {ok, Pid} ->
             NominatedID = Id == 7,
+            ReplPid = 
+                case riakc_pb_socket:start_link(ReplTargetIp, ReplTargetPort) of
+                    {ok, RP} ->
+                        RP;
+                    _ ->
+                        lager:info("Starting with no repl check"),
+                        no_repl_check
+                end,
             {ok, #state {
                pb_pid = Pid,
+               repl_pid = ReplPid,
                http_host = HTTPTargetIp,
                http_port = HTTPTargetPort,
                recordBucket = <<"domainRecord">>,
@@ -306,9 +323,16 @@ run(update_with2i, KeyGen, ValueGen, State) ->
             {error, Reason, State}
     end;
 %% Put an object with a unique key and a non-compressable value
-run(put_unique, _KeyGen, _ValueGen, State) ->
+run(put_unique_bet365, _KeyGen, _ValueGen, State) ->
     Pid = State#state.pb_pid,
-    Bucket = State#state.documentBucket,
+    
+    Bucket = 
+        case erlang:phash2(Pid) rem 2 of
+            0 ->
+               <<"abcdefghijklmnopqrstuvwxyz_1">>;
+            1 ->
+               <<"abcdefghijklmnopqrstuvwxyz_2">>
+        end,
 
     UKC = State#state.unique_key_count,
     Key = 
@@ -319,8 +343,8 @@ run(put_unique, _KeyGen, _ValueGen, State) ->
     Value = non_compressible_value(State#state.unique_size),
 
     Robj0 = riakc_obj:new(Bucket, to_binary(Key)),
-    MD1 = riakc_obj:get_update_metadata(Robj0),
-    MD2 = riakc_obj:set_secondary_index(MD1, generate_binary_indexes()),
+    MD2 = riakc_obj:get_update_metadata(Robj0),
+    % MD2 = riakc_obj:set_secondary_index(MD1, generate_binary_indexes()),
     Robj1 = riakc_obj:update_value(Robj0, Value),
     Robj2 = riakc_obj:update_metadata(Robj1, MD2),
 
@@ -331,6 +355,26 @@ run(put_unique, _KeyGen, _ValueGen, State) ->
         {error, Reason} ->
             {error, Reason, State}
     end;
+%% Put an object with a unique key and a non-compressable value
+run(put_unique, _KeyGen, _ValueGen, State) ->
+    {Pid, _Bucket, _Key, Robj, UKC} = prepare_unique_put(State),
+    %% Write the object...
+    case riakc_pb_socket:put(Pid, Robj, State#state.pb_timeout) of
+        ok ->
+            {ok, State#state{unique_key_count = UKC + 1}};
+        {error, Reason} ->
+            {error, Reason, State}
+    end;
+run(put_unique_checkrepl, _KeyGen, _ValueGen, State) ->
+    {Pid, Bucket, Key, Robj, UKC} = prepare_unique_put(State),
+    %% Write the object...
+    case riakc_pb_socket:put(Pid, Robj, State#state.pb_timeout) of
+        ok ->
+            check_repl(State#state.repl_pid, Bucket, to_binary(Key), State#state.pb_timeout),
+            {ok, State#state{unique_key_count = UKC + 1}};
+        {error, Reason} ->
+            {error, Reason, State}
+    end;
 run(get_unique, _KeyGen, _ValueGen, State) ->    
     % Get one of the objects with unique keys
     Pid = State#state.pb_pid,
@@ -514,6 +558,24 @@ run(Other, _, _, _) ->
 %% ====================================================================
 
 
+prepare_unique_put(State) ->
+    Pid = State#state.pb_pid,
+    Bucket = State#state.documentBucket,
+    
+    UKC = State#state.unique_key_count,
+    Key = 
+        generate_uniquekey(UKC, 
+                            State#state.keyid, 
+                            State#state.unique_keyorder),
+    
+    Value = non_compressible_value(State#state.unique_size),
+    
+    Robj0 = riakc_obj:new(Bucket, to_binary(Key)),
+    MD1 = riakc_obj:get_update_metadata(Robj0),
+    MD2 = riakc_obj:set_secondary_index(MD1, generate_binary_indexes()),
+    Robj1 = riakc_obj:update_value(Robj0, Value),
+    Robj2 = riakc_obj:update_metadata(Robj1, MD2),
+    {Pid, Bucket, Key, Robj2, UKC}.
 
 json_get(Url, Timeout) ->
     json_get(Url, Timeout, true).
@@ -561,6 +623,15 @@ ensure_module(Module) ->
             ok
     end.
 
+check_repl(no_repl_check, _B, _K, _TO) ->
+    ok;
+check_repl(ReplPid, Bucket, Key, Timeout) ->
+    case riakc_pb_socket:get(ReplPid, Bucket, Key, Timeout) of
+        {ok, _Obj} ->
+            ok;
+        {error, _} ->
+            check_repl(ReplPid, Bucket, Key, Timeout)
+    end.
 
 %% ====================================================================
 %% Spawned Runners
@@ -727,4 +798,4 @@ convert_tolist(I) when is_integer(I) ->
     list_to_binary(lists:flatten(io_lib:format("~9..0B", [I])));
 convert_tolist(Bin) ->
     <<I:26/integer, _Tail:6/bitstring>> = Bin,
-    convert_tolist(I).
+    convert_tolist(I).