feat(collector): add statistics for estimate key number of partition

src/server/pegasus_server_impl.cpp

@@ -338,6 +338,13 @@ pegasus_server_impl::pegasus_server_impl(dsn::replication::replica *r)
     snprintf(name, 255, "rdb.memtable.memory_usage@%s", str_gpid.c_str());
     _pfc_rdb_memtable_mem_usage.init_app_counter(
         "app.pegasus", name, COUNTER_TYPE_NUMBER, "statistic the memory usage of rocksdb memtable");
+
+    snprintf(name, 255, "rdb.estimate_key_number@%s", str_gpid.c_str());
+    _pfc_rdb_estimate_key_number.init_app_counter(
+        "app.pegasus",
+        name,
+        COUNTER_TYPE_NUMBER,
+        "statistic the estimate key number of rocksdb all data");
 }
 
 void pegasus_server_impl::parse_checkpoints()
@@ -2242,6 +2249,12 @@ void pegasus_server_impl::update_replica_rocksdb_statistics()
         _pfc_rdb_memtable_mem_usage->set(val);
         dinfo_replica("_pfc_rdb_memtable_mem_usage: {} bytes", val);
     }
+
+    if (_db->GetProperty(rocksdb::DB::Properties::kEstimateNumKeys, &str_val) &&
+        dsn::buf2uint64(str_val, val)) {
+        _pfc_rdb_estimate_key_number->set(val);
+        dinfo_replica("_pfc_rdb_estimate_key_number: {}", val);
+    }
 }
 
 void pegasus_server_impl::update_server_rocksdb_statistics()

src/server/pegasus_server_impl.h

@@ -364,6 +364,7 @@ class pegasus_server_impl : public ::dsn::apps::rrdb_service
     ::dsn::perf_counter_wrapper _pfc_rdb_block_cache_total_count;
     ::dsn::perf_counter_wrapper _pfc_rdb_index_and_filter_blocks_mem_usage;
     ::dsn::perf_counter_wrapper _pfc_rdb_memtable_mem_usage;
+    ::dsn::perf_counter_wrapper _pfc_rdb_estimate_key_number;
 };
 
 } // namespace server

src/server/test/pegasus_server_impl_test.cpp

@@ -3,7 +3,11 @@
 // can be found in the LICENSE file in the root directory of this source tree.
 
 #include <base/pegasus_key_schema.h>
+#include <dsn/utility/defer.h>
+#include "server/pegasus_server_write.h"
+#include "server/pegasus_write_service_impl.h"
 #include "pegasus_server_test_base.h"
+#include "message_utils.h"
 
 namespace pegasus {
 namespace server {
@@ -13,6 +17,8 @@ class pegasus_server_impl_test : public pegasus_server_test_base
 public:
     pegasus_server_impl_test() : pegasus_server_test_base() { start(); }
 
+    rocksdb::DB *get_server_db() { return _server->_db; }
+
     void test_table_level_slow_query()
     {
         // the on_get function will sleep 10ms for unit test,
@@ -59,5 +65,28 @@ class pegasus_server_impl_test : public pegasus_server_test_base
 
 TEST_F(pegasus_server_impl_test, test_table_level_slow_query) { test_table_level_slow_query(); }
 
+TEST_F(pegasus_server_impl_test, test_rdb_estimate_num_keys)
+{
+    std::unique_ptr<pegasus_server_write> server_write =
+        dsn::make_unique<pegasus_server_write>(_server.get(), true);
+
+    dsn::apps::update_request req;
+    pegasus_generate_key(req.key, std::string("hash"), std::string("sort"));
+    req.value.assign("value", 0, 5);
+
+    const int put_cnt = 5;
+    auto writes = new dsn::message_ex *[put_cnt];
+    // generate same five kv, and the kEstimateNumKeys will be count five time
+    for (int i = 0; i < put_cnt; i++) {
+        writes[i] = pegasus::create_put_request(req);
+    }
+    server_write->on_batched_write_requests(writes, put_cnt, 0, 0);
+    auto cleanup = dsn::defer([=]() { delete[] writes; });
+
+    std::string str_val;
+    get_server_db()->GetProperty(rocksdb::DB::Properties::kEstimateNumKeys, &str_val);
+    ASSERT_EQ(atoi(str_val.c_str()), put_cnt);
+}
+
 } // namespace server
 } // namespace pegasus