From a92194e5b2cd555a2318e74074db5583bdfd7678 Mon Sep 17 00:00:00 2001
From: Haobo Xu <haobo@fb.com>
Date: Thu, 27 Mar 2014 11:59:37 -0700
Subject: [PATCH] [RocksDB] Add db property "rocksdb.cur-size-active-mem-table"

Summary: as title

Test Plan: db_test

Reviewers: sdong

Reviewed By: sdong

CC: leveldb

Differential Revision: https://reviews.facebook.net/D17217
---
 db/db_impl.cc        |  3 +--
 db/db_impl.h         |  1 +
 db/db_test.cc        |  7 +++++++
 db/internal_stats.cc | 14 ++++++++++----
 db/internal_stats.h  |  5 +++--
 5 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index 2c24df79d..67daecb48 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -4089,8 +4089,7 @@ bool DBImpl::GetProperty(const Slice& property, std::string* value) {
   value->clear();
   DBPropertyType property_type = GetPropertyType(property);
   MutexLock l(&mutex_);
-  return internal_stats_.GetProperty(property_type, property, value,
-                                     versions_.get(), imm_);
+  return internal_stats_.GetProperty(property_type, property, value, this);
 }
 
 void DBImpl::GetApproximateSizes(
diff --git a/db/db_impl.h b/db/db_impl.h
index 6165f93d3..4cfb6ecaf 100644
--- a/db/db_impl.h
+++ b/db/db_impl.h
@@ -297,6 +297,7 @@ class DBImpl : public DB {
 
  private:
   friend class DB;
+  friend class InternalStats;
   friend class TailingIterator;
   friend struct SuperVersion;
   struct CompactionState;
diff --git a/db/db_test.cc b/db/db_test.cc
index ef39a014b..ee70871a2 100644
--- a/db/db_test.cc
+++ b/db/db_test.cc
@@ -2036,6 +2036,8 @@ TEST(DBTest, NumImmutableMemTable) {
     ASSERT_EQ(1, (int) perf_context.get_from_memtable_count);
 
     ASSERT_OK(dbfull()->Put(writeOpt, "k3", big_value));
+    ASSERT_TRUE(dbfull()->GetProperty("rocksdb.cur-size-active-mem-table",
+                                      &num));
     ASSERT_TRUE(dbfull()->GetProperty("rocksdb.num-immutable-mem-table", &num));
     ASSERT_EQ(num, "2");
     perf_context.Reset();
@@ -2051,6 +2053,11 @@ TEST(DBTest, NumImmutableMemTable) {
     dbfull()->Flush(FlushOptions());
     ASSERT_TRUE(dbfull()->GetProperty("rocksdb.num-immutable-mem-table", &num));
     ASSERT_EQ(num, "0");
+    ASSERT_TRUE(dbfull()->GetProperty("rocksdb.cur-size-active-mem-table",
+                                      &num));
+    // "208" is the size of the metadata of an empty skiplist, this would
+    // break if we change the default skiplist implementation
+    ASSERT_EQ(num, "208");
     SetPerfLevel(kDisable);
   } while (ChangeCompactOptions());
 }
diff --git a/db/internal_stats.cc b/db/internal_stats.cc
index 629941c88..4cc049965 100644
--- a/db/internal_stats.cc
+++ b/db/internal_stats.cc
@@ -7,6 +7,7 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 
 #include "db/internal_stats.h"
+#include "db/db_impl.h"
 #include "db/memtable_list.h"
 
 #include <vector>
@@ -35,15 +36,18 @@ DBPropertyType GetPropertyType(const Slice& property) {
     return kCompactionPending;
   } else if (in == "background-errors") {
     return kBackgroundErrors;
+  } else if (in == "cur-size-active-mem-table") {
+    return kCurSizeActiveMemTable;
   }
   return kUnknown;
 }
 
 bool InternalStats::GetProperty(DBPropertyType property_type,
                                 const Slice& property, std::string* value,
-                                VersionSet* version_set,
-                                const MemTableList& imm) {
+                                DBImpl* db) {
+  VersionSet* version_set = db->versions_.get();
   Version* current = version_set->current();
+  const MemTableList& imm = db->imm_;
   Slice in = property;
 
   switch (property_type) {
@@ -341,12 +345,14 @@ bool InternalStats::GetProperty(DBPropertyType property_type,
       // 0 otherwise,
       *value = std::to_string(current->NeedsCompaction() ? 1 : 0);
       return true;
-    /////////////
     case kBackgroundErrors:
       // Accumulated number of  errors in background flushes or compactions.
       *value = std::to_string(GetBackgroundErrorCount());
       return true;
-    /////////
+    case kCurSizeActiveMemTable:
+      // Current size of the active memtable
+      *value = std::to_string(db->mem_->ApproximateMemoryUsage());
+      return true;
     default:
       return false;
   }
diff --git a/db/internal_stats.h b/db/internal_stats.h
index b6032d014..e140e7280 100644
--- a/db/internal_stats.h
+++ b/db/internal_stats.h
@@ -19,6 +19,7 @@
 namespace rocksdb {
 
 class MemTableList;
+class DBImpl;
 
 enum DBPropertyType {
   kNumFilesAtLevel,  // Number of files at a specific level
@@ -31,6 +32,7 @@ enum DBPropertyType {
                           // 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
   kUnknown,
 };
 
@@ -124,8 +126,7 @@ class InternalStats {
   uint64_t BumpAndGetBackgroundErrorCount() { return ++bg_error_count_; }
 
   bool GetProperty(DBPropertyType property_type, const Slice& property,
-                   std::string* value, VersionSet* version_set,
-                   const MemTableList& imm);
+                   std::string* value, DBImpl* db);
 
  private:
   std::vector<CompactionStats> compaction_stats_;