Expose number of entries in mem tables to users

Summary: In this patch, two new DB properties are defined: rocksdb.num-immutable-mem-table and rocksdb.num-entries-imm-mem-tables, from where number of entries in mem tables can be exposed to users

Test Plan:
Cover the codes in db_test
make all check

Reviewers: haobo, ljin, igor

Reviewed By: igor

CC: nkg-, igor, yhchiang, dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D18207

db/db_test.cc

@@ -1009,12 +1009,28 @@ TEST(DBTest, Empty) {
     options.write_buffer_size = 100000;  // Small write buffer
     CreateAndReopenWithCF({"pikachu"}, &options);
 
+    std::string num;
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ("0", num);
+
     ASSERT_OK(Put(1, "foo", "v1"));
     ASSERT_EQ("v1", Get(1, "foo"));
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ("1", num);
 
     env_->delay_sstable_sync_.Release_Store(env_);  // Block sync calls
     Put(1, "k1", std::string(100000, 'x'));         // Fill memtable
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ("2", num);
+
     Put(1, "k2", std::string(100000, 'y'));         // Trigger compaction
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ("1", num);
+
     ASSERT_EQ("v1", Get(1, "foo"));
     env_->delay_sstable_sync_.Release_Store(nullptr);   // Release sync calls
   } while (ChangeOptions());
@@ -2225,6 +2241,9 @@ TEST(DBTest, NumImmutableMemTable) {
     ASSERT_TRUE(dbfull()->GetProperty(handles_[1],
                                       "rocksdb.num-immutable-mem-table", &num));
     ASSERT_EQ(num, "0");
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ(num, "1");
     perf_context.Reset();
     Get(1, "k1");
     ASSERT_EQ(1, (int) perf_context.get_from_memtable_count);
@@ -2233,6 +2252,13 @@ TEST(DBTest, NumImmutableMemTable) {
     ASSERT_TRUE(dbfull()->GetProperty(handles_[1],
                                       "rocksdb.num-immutable-mem-table", &num));
     ASSERT_EQ(num, "1");
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ(num, "1");
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-imm-mem-tables", &num));
+    ASSERT_EQ(num, "1");
+
     perf_context.Reset();
     Get(1, "k1");
     ASSERT_EQ(2, (int) perf_context.get_from_memtable_count);
@@ -2246,6 +2272,12 @@ TEST(DBTest, NumImmutableMemTable) {
     ASSERT_TRUE(dbfull()->GetProperty(handles_[1],
                                       "rocksdb.num-immutable-mem-table", &num));
     ASSERT_EQ(num, "2");
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-active-mem-table", &num));
+    ASSERT_EQ(num, "1");
+    ASSERT_TRUE(dbfull()->GetProperty(
+        handles_[1], "rocksdb.num-entries-imm-mem-tables", &num));
+    ASSERT_EQ(num, "2");
     perf_context.Reset();
     Get(1, "k2");
     ASSERT_EQ(2, (int) perf_context.get_from_memtable_count);

db/internal_stats.cc

@@ -37,6 +37,10 @@ DBPropertyType GetPropertyType(const Slice& property) {
     return kBackgroundErrors;
   } else if (in == "cur-size-active-mem-table") {
     return kCurSizeActiveMemTable;
+  } else if (in == "num-entries-active-mem-table") {
+    return kNumEntriesInMutableMemtable;
+  } else if (in == "num-entries-imm-mem-tables") {
+    return kNumEntriesInImmutableMemtable;
   }
   return kUnknown;
 }
@@ -349,6 +353,14 @@ bool InternalStats::GetProperty(DBPropertyType property_type,
       // Current size of the active memtable
       *value = std::to_string(cfd->mem()->ApproximateMemoryUsage());
       return true;
+    case kNumEntriesInMutableMemtable:
+      // Current size of the active memtable
+      *value = std::to_string(cfd->mem()->GetNumEntries());
+      return true;
+    case kNumEntriesInImmutableMemtable:
+      // Current size of the active memtable
+      *value = std::to_string(cfd->imm()->current()->GetTotalNumEntries());
+      return true;
     default:
       return false;
   }

db/internal_stats.h

@@ -28,13 +28,16 @@ enum DBPropertyType {
   kLevelStats,       // Return number of files and total sizes of each level
   kStats,            // Return general statitistics of DB
   kSsTables,         // Return a human readable string of current SST files
-  kNumImmutableMemTable,  // Return number of immutable mem tables
-  kMemtableFlushPending,  // Return 1 if mem table flushing is pending,
-                          // otherwise
-                          // 0.
-  kCompactionPending,     // Return 1 if a compaction is pending. Otherwise 0.
-  kBackgroundErrors,      // Return accumulated background errors encountered.
+  kNumImmutableMemTable,   // Return number of immutable mem tables
+  kMemtableFlushPending,   // Return 1 if mem table flushing is pending,
+                           // otherwise 0.
+  kCompactionPending,      // Return 1 if a compaction is pending. Otherwise 0.
+  kBackgroundErrors,       // Return accumulated background errors encountered.
   kCurSizeActiveMemTable,  // Return current size of the active memtable
+  kNumEntriesInMutableMemtable,    // Return number of entries in the mutable
+                                   // memtable.
+  kNumEntriesInImmutableMemtable,  // Return sum of number of entries in all
+                                   // the immutable mem tables.
   kUnknown,
 };
 

db/memtable.cc

@@ -29,8 +29,7 @@
 
 namespace rocksdb {
 
-MemTable::MemTable(const InternalKeyComparator& cmp,
-                   const Options& options)
+MemTable::MemTable(const InternalKeyComparator& cmp, const Options& options)
     : comparator_(cmp),
       refs_(0),
       kArenaBlockSize(OptimizeBlockSize(options.arena_block_size)),
@@ -38,6 +37,7 @@ MemTable::MemTable(const InternalKeyComparator& cmp,
       arena_(options.arena_block_size),
       table_(options.memtable_factory->CreateMemTableRep(
           comparator_, &arena_, options.prefix_extractor.get())),
+      num_entries_(0),
       flush_in_progress_(false),
       flush_completed_(false),
       file_number_(0),
@@ -259,6 +259,7 @@ void MemTable::Add(SequenceNumber s, ValueType type,
   memcpy(p, value.data(), val_size);
   assert((unsigned)(p + val_size - buf) == (unsigned)encoded_len);
   table_->Insert(handle);
+  num_entries_++;
 
   if (prefix_bloom_) {
     assert(prefix_extractor_);

db/memtable.h

@@ -132,6 +132,9 @@ class MemTable {
   // key in the memtable.
   size_t CountSuccessiveMergeEntries(const LookupKey& key);
 
+  // Get total number of entries in the mem table.
+  uint64_t GetNumEntries() const { return num_entries_; }
+
   // Returns the edits area that is needed for flushing the memtable
   VersionEdit* GetEdits() { return &edit_; }
 
@@ -174,6 +177,8 @@ class MemTable {
   Arena arena_;
   unique_ptr<MemTableRep> table_;
 
+  uint64_t num_entries_;
+
   // These are used to manage memtable flushes to storage
   bool flush_in_progress_; // started the flush
   bool flush_completed_;   // finished the flush

db/memtable_list.cc

@@ -78,6 +78,14 @@ void MemTableListVersion::AddIterators(const ReadOptions& options,
   }
 }
 
+uint64_t MemTableListVersion::GetTotalNumEntries() const {
+  uint64_t total_num = 0;
+  for (auto& m : memlist_) {
+    total_num += m->GetNumEntries();
+  }
+  return total_num;
+}
+
 // caller is responsible for referencing m
 void MemTableListVersion::Add(MemTable* m) {
   assert(refs_ == 1);  // only when refs_ == 1 is MemTableListVersion mutable

db/memtable_list.h

@@ -49,6 +49,8 @@ class MemTableListVersion {
   void AddIterators(const ReadOptions& options,
                     std::vector<Iterator*>* iterator_list);
 
+  uint64_t GetTotalNumEntries() const;
+
  private:
   // REQUIRE: m is mutable memtable
   void Add(MemTable* m);