From e3004faa6e3c8be3093be36f08d3e5992b8f7d0b Mon Sep 17 00:00:00 2001
From: martinsumner <martin.sumner@adaptip.co.uk>
Date: Thu, 17 Nov 2016 15:55:29 +0000
Subject: [PATCH] Add support for folding keys on a single bucket

---
 src/leveled_bookie.erl                            | 13 ++++++++++---
 test/end_to_end/testutil.erl                      |  5 +++++
 test/volume/src/basho_bench_driver_eleveleddb.erl | 12 ++++++------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/leveled_bookie.erl b/src/leveled_bookie.erl
index 06ba1c9..2ebc012 100644
--- a/src/leveled_bookie.erl
+++ b/src/leveled_bookie.erl
@@ -349,6 +349,10 @@ handle_call({return_folder, FolderType}, _From, State) ->
             {reply,
                 allkey_query(State, Tag),
                 State};
+        {keylist, Tag, Bucket} ->
+            {reply,
+                bucketkey_query(State, Tag, Bucket),
+                State};
         {hashtree_query, Tag, JournalCheck} ->
             {reply,
                 hashtree_query(State, Tag, JournalCheck),
@@ -535,7 +539,7 @@ foldobjects(State, Tag, StartKey, EndKey, FoldObjectsFun) ->
     {async, Folder}.
 
 
-allkey_query(State, Tag) ->
+bucketkey_query(State, Tag, Bucket) ->
     {ok,
         {LedgerSnapshot, LedgerCache},
         _JournalSnapshot} = snapshot_store(State, ledger),
@@ -543,8 +547,8 @@ allkey_query(State, Tag) ->
                 leveled_log:log("B0004", [gb_trees:size(LedgerCache)]),
                 ok = leveled_penciller:pcl_loadsnapshot(LedgerSnapshot,
                                                             LedgerCache),
-                SK = leveled_codec:to_ledgerkey(null, null, Tag),
-                EK = leveled_codec:to_ledgerkey(null, null, Tag),
+                SK = leveled_codec:to_ledgerkey(Bucket, null, Tag),
+                EK = leveled_codec:to_ledgerkey(Bucket, null, Tag),
                 AccFun = accumulate_keys(),
                 Acc = leveled_penciller:pcl_fetchkeys(LedgerSnapshot,
                                                         SK,
@@ -556,6 +560,9 @@ allkey_query(State, Tag) ->
                 end,
     {async, Folder}.
 
+allkey_query(State, Tag) ->
+    bucketkey_query(State, Tag, null).
+
 
 snapshot_store(State, SnapType) ->
     PCLopts = #penciller_options{start_snapshot=true,
diff --git a/test/end_to_end/testutil.erl b/test/end_to_end/testutil.erl
index 61aa387..e08de45 100644
--- a/test/end_to_end/testutil.erl
+++ b/test/end_to_end/testutil.erl
@@ -428,6 +428,11 @@ rotating_object_check(RootPath, B, NumberOfObjects) ->
     ok = testutil:check_indexed_objects(Book2, B, KSpcL3, V3),
     {KSpcL4, V4} = testutil:put_altered_indexed_objects(Book2, B, KSpcL3),
     ok = testutil:check_indexed_objects(Book2, B, KSpcL4, V4),
+    {async, BList} = leveled_bookie:book_returnfolder(Book2,
+                                                        {keylist,
+                                                            ?RIAK_TAG,
+                                                            B}),
+    true = NumberOfObjects == length(BList()),
     ok = leveled_bookie:book_close(Book2),
     ok.
     
diff --git a/test/volume/src/basho_bench_driver_eleveleddb.erl b/test/volume/src/basho_bench_driver_eleveleddb.erl
index 9aedcf7..ea81e90 100644
--- a/test/volume/src/basho_bench_driver_eleveleddb.erl
+++ b/test/volume/src/basho_bench_driver_eleveleddb.erl
@@ -27,7 +27,7 @@
 -export([new/1,
          run/4]).
 
--include("basho_bench.hrl").
+% -include("basho_bench.hrl").
 
 -record(state, {
           instance
@@ -37,7 +37,7 @@ get_instances() ->
     case basho_bench_config:get(eleveleddb_instances, undefined) of
         undefined ->
             Instances = start_instances(),
-            ?INFO("Instances started ~w~n", [Instances]),
+            % ?INFO("Instances started ~w~n", [Instances]),
             basho_bench_config:set(eleveleddb_instances, Instances),
             Instances;
         Instances ->
@@ -48,11 +48,11 @@ get_instances() ->
 start_instances() ->
     BaseDir = basho_bench_config:get(eleveleddb_dir, "."),
     Num = basho_bench_config:get(eleveleddb_num_instances, 1),
-    ?INFO("Starting up ~p eleveleddb instances under ~s .\n",
-          [Num, BaseDir]),
+    % ?INFO("Starting up ~p eleveleddb instances under ~s .\n",
+    %       [Num, BaseDir]),
     Refs = [begin
                 Dir = filename:join(BaseDir, "instance." ++ integer_to_list(N)),
-                ?INFO("Opening eleveleddb instance in ~s\n", [Dir]),
+                % ?INFO("Opening eleveleddb instance in ~s\n", [Dir]),
                 {ok, Ref} = leveled_bookie:book_start(Dir, 2000, 500000000),
                 Ref
             end || N <- lists:seq(1, Num)],
@@ -62,7 +62,7 @@ new(Id) ->
     Instances = get_instances(),
     Count = size(Instances),
     Idx = ((Id - 1) rem Count) + 1,
-    ?INFO("Worker ~p using instance ~p.\n", [Id, Idx]),
+    % ?INFO("Worker ~p using instance ~p.\n", [Id, Idx]),
     State = #state{instance = element(Idx, Instances)},
     {ok, State}.