From 528c2bc6ad03875865be8712aacdf319f7002d44 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 3fd21b047f7d89..dc645d041eaf40 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 74ba8b05e86713..a0b08bc19c6510 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 5f86cd03f30261..6819d5bc49f674 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 4c169bf22ed92e..53c7068765e2d9 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 7f5fc07e42337b..0881bce9c53fa5 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