From 36fc95597178cd694fa4775a773ef8da8337bf08 Mon Sep 17 00:00:00 2001
From: ssid <ssid@google.com>
Date: Tue, 29 Sep 2015 11:52:21 -0700
Subject: [PATCH] Add "approximate-memory-usage" property to
 leveldb::DB::GetProperty

The approximate RAM usage of the database is calculated from the memory
allocated for write buffers and the block cache. This is to give an
estimate of memory usage to leveldb clients.
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=104222307
---
 db/db_impl.cc           | 13 +++++++++++++
 db/db_test.cc           | 11 +++++++++++
 include/leveldb/cache.h |  4 ++++
 include/leveldb/db.h    |  2 ++
 util/cache.cc           | 13 ++++++++++++-
 5 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index 3fd21b0..dc645d0 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -1425,6 +1425,19 @@ bool DBImpl::GetProperty(const Slice& property, std::string* value) {
   } else if (in == "sstables") {
     *value = versions_->current()->DebugString();
     return true;
+  } else if (in == "approximate-memory-usage") {
+    size_t total_usage = options_.block_cache->TotalCharge();
+    if (mem_) {
+      total_usage += mem_->ApproximateMemoryUsage();
+    }
+    if (imm_) {
+      total_usage += imm_->ApproximateMemoryUsage();
+    }
+    char buf[50];
+    snprintf(buf, sizeof(buf), "%llu",
+             static_cast<unsigned long long>(total_usage));
+    value->append(buf);
+    return true;
   }
 
   return false;
diff --git a/db/db_test.cc b/db/db_test.cc
index 74ba8b0..a0b08bc 100644
--- a/db/db_test.cc
+++ b/db/db_test.cc
@@ -563,6 +563,17 @@ TEST(DBTest, GetFromVersions) {
   } while (ChangeOptions());
 }
 
+TEST(DBTest, GetMemUsage) {
+  do {
+    ASSERT_OK(Put("foo", "v1"));
+    std::string val;
+    ASSERT_TRUE(db_->GetProperty("leveldb.approximate-memory-usage", &val));
+    int mem_usage = atoi(val.c_str());
+    ASSERT_GT(mem_usage, 0);
+    ASSERT_LT(mem_usage, 5*1024*1024);
+  } while (ChangeOptions());
+}
+
 TEST(DBTest, GetSnapshot) {
   do {
     // Try with both a short key and a long key
diff --git a/include/leveldb/cache.h b/include/leveldb/cache.h
index 5f86cd0..6819d5b 100644
--- a/include/leveldb/cache.h
+++ b/include/leveldb/cache.h
@@ -88,6 +88,10 @@ class Cache {
   // leveldb may change Prune() to a pure abstract method.
   virtual void Prune() {}
 
+  // Return an estimate of the combined charges of all elements stored in the
+  // cache.
+  virtual size_t TotalCharge() const = 0;
+
  private:
   void LRU_Remove(Handle* e);
   void LRU_Append(Handle* e);
diff --git a/include/leveldb/db.h b/include/leveldb/db.h
index 4c169bf..53c7068 100644
--- a/include/leveldb/db.h
+++ b/include/leveldb/db.h
@@ -115,6 +115,8 @@ class DB {
   //     about the internal operation of the DB.
   //  "leveldb.sstables" - returns a multi-line string that describes all
   //     of the sstables that make up the db contents.
+  //  "leveldb.approximate-memory-usage" - returns the approximate number of
+  //     bytes of memory in use by the DB.
   virtual bool GetProperty(const Slice& property, std::string* value) = 0;
 
   // For each i in [0,n-1], store in "sizes[i]", the approximate
diff --git a/util/cache.cc b/util/cache.cc
index 7f5fc07..0881bce 100644
--- a/util/cache.cc
+++ b/util/cache.cc
@@ -148,6 +148,10 @@ class LRUCache {
   void Release(Cache::Handle* handle);
   void Erase(const Slice& key, uint32_t hash);
   void Prune();
+  size_t TotalCharge() const {
+    MutexLock l(&mutex_);
+    return usage_;
+  }
 
  private:
   void LRU_Remove(LRUHandle* e);
@@ -158,7 +162,7 @@ class LRUCache {
   size_t capacity_;
 
   // mutex_ protects the following state.
-  port::Mutex mutex_;
+  mutable port::Mutex mutex_;
   size_t usage_;
 
   // Dummy head of LRU list.
@@ -333,6 +337,13 @@ class ShardedLRUCache : public Cache {
       shard_[s].Prune();
     }
   }
+  virtual size_t TotalCharge() const {
+    size_t total = 0;
+    for (int s = 0; s < kNumShards; s++) {
+      total += shard_[s].TotalCharge();
+    }
+    return total;
+  }
 };
 
 }  // end anonymous namespace