From 1543d5d92e5d1b144b75dde6d76132dea7a1a20f Mon Sep 17 00:00:00 2001
From: Yi Wu <yiwu@fb.com>
Date: Tue, 15 Nov 2016 20:05:36 -0800
Subject: [PATCH] Report memory usage by memtable insert hints map.

Summary:
It is hard to measure acutal memory usage by std containers. Even
providing a custom allocator will miss count some of the usage. Here we
only do a wild guess on its memory usage.
Closes https://github.com/facebook/rocksdb/pull/1511

Differential Revision: D4179945

Pulled By: yiwu-arbug

fbshipit-source-id: 32ab929
---
 db/db_properties_test.cc |  7 ++++---
 db/memtable.cc           | 25 ++++++++++++++++---------
 util/autovector.h        | 14 ++++++++++++--
 util/memory_usage.h      | 25 +++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 14 deletions(-)
 create mode 100644 util/memory_usage.h

diff --git a/db/db_properties_test.cc b/db/db_properties_test.cc
index e0ad3ab9f..31afc2cfd 100644
--- a/db/db_properties_test.cc
+++ b/db/db_properties_test.cc
@@ -491,6 +491,7 @@ TEST_F(DBPropertiesTest, NumImmutableMemTable) {
 
     std::string big_value(1000000 * 2, 'x');
     std::string num;
+    uint64_t value;
     SetPerfLevel(kEnableTime);
     ASSERT_TRUE(GetPerfLevel() == kEnableTime);
 
@@ -555,11 +556,11 @@ TEST_F(DBPropertiesTest, NumImmutableMemTable) {
     ASSERT_TRUE(dbfull()->GetProperty(
         handles_[1], DB::Properties::kNumImmutableMemTableFlushed, &num));
     ASSERT_EQ(num, "3");
-    ASSERT_TRUE(dbfull()->GetProperty(
-        handles_[1], "rocksdb.cur-size-active-mem-table", &num));
+    ASSERT_TRUE(dbfull()->GetIntProperty(
+        handles_[1], "rocksdb.cur-size-active-mem-table", &value));
     // "384" is the size of the metadata of two empty skiplists, this would
     // break if we change the default skiplist implementation
-    ASSERT_EQ(num, "384");
+    ASSERT_GE(value, 384);
 
     uint64_t int_num;
     uint64_t base_total_size;
diff --git a/db/memtable.cc b/db/memtable.cc
index f12365298..9e059409f 100644
--- a/db/memtable.cc
+++ b/db/memtable.cc
@@ -17,6 +17,7 @@
 #include "db/merge_context.h"
 #include "db/merge_helper.h"
 #include "db/pinned_iterators_manager.h"
+#include "port/port.h"
 #include "rocksdb/comparator.h"
 #include "rocksdb/env.h"
 #include "rocksdb/iterator.h"
@@ -27,7 +28,9 @@
 #include "table/iterator_wrapper.h"
 #include "table/merger.h"
 #include "util/arena.h"
+#include "util/autovector.h"
 #include "util/coding.h"
+#include "util/memory_usage.h"
 #include "util/murmurhash.h"
 #include "util/mutexlock.h"
 #include "util/perf_context_imp.h"
@@ -105,17 +108,21 @@ MemTable::MemTable(const InternalKeyComparator& cmp,
 MemTable::~MemTable() { assert(refs_ == 0); }
 
 size_t MemTable::ApproximateMemoryUsage() {
-  size_t arena_usage = arena_.ApproximateMemoryUsage();
-  size_t table_usage = table_->ApproximateMemoryUsage();
-  table_usage += range_del_table_->ApproximateMemoryUsage();
-  // let MAX_USAGE =  std::numeric_limits<size_t>::max()
-  // then if arena_usage + total_usage >= MAX_USAGE, return MAX_USAGE.
-  // the following variation is to avoid numeric overflow.
-  if (arena_usage >= std::numeric_limits<size_t>::max() - table_usage) {
-    return std::numeric_limits<size_t>::max();
+  autovector<size_t> usages = {arena_.ApproximateMemoryUsage(),
+                               table_->ApproximateMemoryUsage(),
+                               range_del_table_->ApproximateMemoryUsage(),
+                               rocksdb::ApproximateMemoryUsage(insert_hints_)};
+  size_t total_usage = 0;
+  for (size_t usage : usages) {
+    // If usage + total_usage >= kMaxSizet, return kMaxSizet.
+    // the following variation is to avoid numeric overflow.
+    if (usage >= port::kMaxSizet - total_usage) {
+      return port::kMaxSizet;
+    }
+    total_usage += usage;
   }
   // otherwise, return the actual usage
-  return arena_usage + table_usage;
+  return total_usage;
 }
 
 bool MemTable::ShouldFlushNow() const {
diff --git a/util/autovector.h b/util/autovector.h
index 74fcc70ec..cfe703ed9 100644
--- a/util/autovector.h
+++ b/util/autovector.h
@@ -6,15 +6,18 @@
 
 #include <algorithm>
 #include <cassert>
-#include <stdexcept>
+#include <initializer_list>
 #include <iterator>
+#include <stdexcept>
 #include <vector>
 
 namespace rocksdb {
 
 #ifdef ROCKSDB_LITE
 template <class T, size_t kSize = 8>
-class autovector : public std::vector<T> {};
+class autovector : public std::vector<T> {
+  using std::vector<T>::vector;
+};
 #else
 // A vector that leverages pre-allocated stack-based array to achieve better
 // performance for array with small amount of items.
@@ -165,6 +168,13 @@ class autovector {
   typedef std::reverse_iterator<const_iterator> const_reverse_iterator;
 
   autovector() = default;
+
+  autovector(std::initializer_list<T> init_list) {
+    for (const T& item : init_list) {
+      push_back(item);
+    }
+  }
+
   ~autovector() = default;
 
   // -- Immutable operations
diff --git a/util/memory_usage.h b/util/memory_usage.h
new file mode 100644
index 000000000..cfe0025fe
--- /dev/null
+++ b/util/memory_usage.h
@@ -0,0 +1,25 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under the BSD-style license found in the
+//  LICENSE file in the root directory of this source tree. An additional grant
+//  of patent rights can be found in the PATENTS file in the same directory.
+
+#pragma once
+
+#include <unordered_map>
+
+namespace rocksdb {
+
+// Helper methods to estimate memroy usage by std containers.
+
+template <class Key, class Value, class Hash>
+size_t ApproximateMemoryUsage(
+    const std::unordered_map<Key, Value, Hash>& umap) {
+  typedef std::unordered_map<Key, Value, Hash> Map;
+  return sizeof(umap) +
+         // Size of all items plus a next pointer for each item.
+         (sizeof(typename Map::value_type) + sizeof(void*)) * umap.size() +
+         // Size of hash buckets.
+         umap.bucket_count() * sizeof(void*);
+}
+
+}  // namespace rocksdb