Implement `couch_db:get_partition_info/2`

This feature allows us to fetch statistics for a given partition key
which will allow for users to find bloated partitions and such forth.

Co-authored-by: Garren Smith <garren.smith@gmail.com>
Co-authored-by: Robert Newson <rnewson@apache.org>

4 files changed, 86 insertions, 0 deletions

diff --git a/src/couch/src/couch_bt_engine.erl b/src/couch/src/couch_bt_engine.erl
index 946b74d0c..7b33c4203 100644
--- a/src/couch/src/couch_bt_engine.erl
+++ b/src/couch/src/couch_bt_engine.erl
@@ -42,6 +42,7 @@
     get_security/1,
     get_props/1,
     get_size_info/1,
+    get_partition_info/2,
     get_update_seq/1,
     get_uuid/1,
 
@@ -277,6 +278,48 @@ get_size_info(#st{} = St) ->
     ].
 
 
+partition_size_cb(traverse, Key, {DC, DDC, Sizes}, {Partition, DCAcc, DDCAcc, SizesAcc}) ->
+    case couch_partition:is_member(Key, Partition) of
+        true ->
+            {skip, {Partition, DC + DCAcc, DDC + DDCAcc, reduce_sizes(Sizes, SizesAcc)}};
+        false ->
+            {ok, {Partition, DCAcc, DDCAcc, SizesAcc}}
+    end;
+
+partition_size_cb(visit, FDI, _PrevReds, {Partition, DCAcc, DDCAcc, Acc}) ->
+    InPartition = couch_partition:is_member(FDI#full_doc_info.id, Partition),
+    Deleted = FDI#full_doc_info.deleted,
+    case {InPartition, Deleted} of
+        {true, true} ->
+            {ok, {Partition, DCAcc, DDCAcc + 1,
+                reduce_sizes(FDI#full_doc_info.sizes, Acc)}};
+        {true, false} ->
+            {ok, {Partition, DCAcc + 1, DDCAcc,
+                reduce_sizes(FDI#full_doc_info.sizes, Acc)}};
+        {false, _} ->
+            {ok, {Partition, DCAcc, DDCAcc, Acc}}
+    end.
+
+
+get_partition_info(#st{} = St, Partition) ->
+    StartKey = couch_partition:start_key(Partition),
+    EndKey = couch_partition:end_key(Partition),
+    Fun = fun partition_size_cb/4,
+    InitAcc = {Partition, 0, 0, #size_info{}},
+    Options = [{start_key, StartKey}, {end_key, EndKey}],
+    {ok, _, OutAcc} = couch_btree:fold(St#st.id_tree, Fun, InitAcc, Options),
+    {Partition, DocCount, DocDelCount, SizeInfo} = OutAcc,
+    [
+        {partition, Partition},
+        {doc_count, DocCount},
+        {doc_del_count, DocDelCount},
+        {sizes, [
+            {active, SizeInfo#size_info.active},
+            {external, SizeInfo#size_info.external}
+        ]}
+    ].
+
+
 get_security(#st{header = Header} = St) ->
     case couch_bt_engine_header:get(Header, security_ptr) of
         undefined ->
diff --git a/src/couch/src/couch_db.erl b/src/couch/src/couch_db.erl
index 2c6f41bf7..74f4a099d 100644
--- a/src/couch/src/couch_db.erl
+++ b/src/couch/src/couch_db.erl
@@ -38,6 +38,7 @@
     get_compacted_seq/1,
     get_compactor_pid/1,
     get_db_info/1,
+    get_partition_info/2,
     get_del_doc_count/1,
     get_doc_count/1,
     get_epochs/1,
@@ -633,6 +634,13 @@ get_db_info(Db) ->
     ],
     {ok, InfoList}.
 
+get_partition_info(#db{} = Db, Partition) when is_binary(Partition) ->
+    Info = couch_db_engine:get_partition_info(Db, Partition),
+    {ok, Info};
+get_partition_info(_Db, _Partition) ->
+    throw({bad_request, <<"`partition` is not valid">>}).
+
+
 get_design_doc(#db{name = <<"shards/", _/binary>> = ShardDbName}, DDocId0) ->
     DDocId = couch_util:normalize_ddoc_id(DDocId0),
     DbName = mem3:dbname(ShardDbName),
diff --git a/src/couch/src/couch_db_engine.erl b/src/couch/src/couch_db_engine.erl
index 806d352cb..91d35b0c7 100644
--- a/src/couch/src/couch_db_engine.erl
+++ b/src/couch/src/couch_db_engine.erl
@@ -44,6 +44,12 @@
 -type purge_info() :: {purge_seq(), uuid(), docid(), revs()}.
 -type epochs() :: [{Node::atom(), UpdateSeq::non_neg_integer()}].
 -type size_info() :: [{Name::atom(), Size::non_neg_integer()}].
+-type partition_info() :: [
+    {partition, Partition::binary()} |
+    {doc_count, DocCount::non_neg_integer()} |
+    {doc_del_count, DocDelCount::non_neg_integer()} |
+    {sizes, size_info()}
+].
 
 -type write_stream_options() :: [
         {buffer_size, Size::pos_integer()} |
@@ -263,6 +269,18 @@
 -callback get_size_info(DbHandle::db_handle()) -> SizeInfo::size_info().
 
 
+% This returns the information for the given partition.
+% It should just be a list of {Name::atom(), Size::non_neg_integer()}
+% It returns the partition name, doc count, deleted doc count and two sizes:
+%
+%   active   - Theoretical minimum number of bytes to store this partition on disk
+%
+%   external - Number of bytes that would be required to represent the
+%              contents of this partition outside of the database
+-callback get_partition_info(DbHandle::db_handle(), Partition::binary()) ->
+    partition_info().
+
+
 % The current update sequence of the database. The update
 % sequence should be incrememnted for every revision added to
 % the database.
@@ -685,6 +703,7 @@
     get_security/1,
     get_props/1,
     get_size_info/1,
+    get_partition_info/2,
     get_update_seq/1,
     get_uuid/1,
 
@@ -861,6 +880,11 @@ get_size_info(#db{} = Db) ->
     Engine:get_size_info(EngineState).
 
 
+get_partition_info(#db{} = Db, Partition) ->
+    #db{engine = {Engine, EngineState}} = Db,
+    Engine:get_partition_info(EngineState, Partition).
+
+
 get_update_seq(#db{} = Db) ->
     #db{engine = {Engine, EngineState}} = Db,
     Engine:get_update_seq(EngineState).
diff --git a/src/couch/src/couch_partition.erl b/src/couch/src/couch_partition.erl
index 783921f0a..9ff77a0ff 100644
--- a/src/couch/src/couch_partition.erl
+++ b/src/couch/src/couch_partition.erl
@@ -18,6 +18,9 @@
     from_docid/1,
     is_member/2,
 
+    start_key/1,
+    end_key/1,
+
     validate_dbname/2,
     validate_docid/1,
     validate_partition/1,
@@ -59,6 +62,14 @@ is_member(DocId, Partition) ->
     end.
 
 
+start_key(Partition) ->
+    <<Partition/binary, ":">>.
+
+
+end_key(Partition) ->
+    <<Partition/binary, ";">>.
+
+
 validate_dbname(DbName, Options) when is_list(DbName) ->
     validate_dbname(?l2b(DbName), Options);
 validate_dbname(DbName, Options) when is_binary(DbName) ->