Use IterKey instead of string in Block::Iter to reduce malloc

Summary:
  Modify a functioin TrimAppend in dbformat.h: IterKey. Write a test for it in dbformat_test
  Use IterKey in block::Iter to replace std::string to reduce malloc.

  Evaluate it using perf record.
  malloc: 4.26% -> 2.91%
  free: 3.61% -> 3.08%

Test Plan:
  make all check
  ./valgrind db_test dbformat_test

Reviewers: ljin, haobo, yhchiang, dhruba, igor, sdong

Reviewed By: sdong

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

db/dbformat.h

@@ -9,6 +9,7 @@
 
 #pragma once
 #include <stdio.h>
+#include <string>
 #include "rocksdb/comparator.h"
 #include "rocksdb/db.h"
 #include "rocksdb/filter_policy.h"
@@ -254,8 +255,38 @@ class IterKey {
 
   Slice GetKey() const { return Slice(key_, key_size_); }
 
+  const size_t Size() { return key_size_; }
+
   void Clear() { key_size_ = 0; }
 
+  // Append "non_shared_data" to its back, from "shared_len"
+  // This function is used in Block::Iter::ParseNextKey
+  // shared_len: bytes in [0, shard_len-1] would be remained
+  // non_shared_data: data to be append, its length must be >= non_shared_len
+  void TrimAppend(const size_t shared_len, const char* non_shared_data,
+                  const size_t non_shared_len) {
+    assert(shared_len <= key_size_);
+
+    size_t total_size = shared_len + non_shared_len;
+    if (total_size <= buf_size_) {
+      key_size_ = total_size;
+    } else {
+      // Need to allocate space, delete previous space
+      char* p = new char[total_size];
+      memcpy(p, key_, shared_len);
+
+      if (key_ != nullptr && key_ != space_) {
+        delete[] key_;
+      }
+
+      key_ = p;
+      key_size_ = total_size;
+      buf_size_ = total_size;
+    }
+
+    memcpy(key_ + shared_len, non_shared_data, non_shared_len);
+  }
+
   void SetKey(const Slice& key) {
     size_t size = key.size();
     EnlargeBufferIfNeeded(size);

db/dbformat_test.cc

@@ -110,6 +110,45 @@ TEST(FormatTest, InternalKeyShortestSuccessor) {
             ShortSuccessor(IKey("\xff\xff", 100, kTypeValue)));
 }
 
+TEST(FormatTest, IterKeyOperation) {
+  IterKey k;
+  const char p[] = "abcdefghijklmnopqrstuvwxyz";
+  const char q[] = "0123456789";
+
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string(""));
+
+  k.TrimAppend(0, p, 3);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abc"));
+
+  k.TrimAppend(1, p, 3);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("aabc"));
+
+  k.TrimAppend(0, p, 26);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abcdefghijklmnopqrstuvwxyz"));
+
+  k.TrimAppend(26, q, 10);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abcdefghijklmnopqrstuvwxyz0123456789"));
+
+  k.TrimAppend(36, q, 1);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abcdefghijklmnopqrstuvwxyz01234567890"));
+
+  k.TrimAppend(26, q, 1);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abcdefghijklmnopqrstuvwxyz0"));
+
+  // Size going up, memory allocation is triggered
+  k.TrimAppend(27, p, 26);
+  ASSERT_EQ(std::string(k.GetKey().data(), k.GetKey().size()),
+            std::string("abcdefghijklmnopqrstuvwxyz0"
+              "abcdefghijklmnopqrstuvwxyz"));
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {

table/block.cc

@@ -22,6 +22,7 @@
 #include "table/format.h"
 #include "util/coding.h"
 #include "util/logging.h"
+#include "db/dbformat.h"
 
 namespace rocksdb {
 
@@ -94,7 +95,7 @@ class Block::Iter : public Iterator {
   // current_ is offset in data_ of current entry.  >= restarts_ if !Valid
   uint32_t current_;
   uint32_t restart_index_;  // Index of restart block in which current_ falls
-  std::string key_;
+  IterKey key_;
   Slice value_;
   Status status_;
   BlockHashIndex* hash_index_;
@@ -115,7 +116,7 @@ class Block::Iter : public Iterator {
   }
 
   void SeekToRestartPoint(uint32_t index) {
-    key_.clear();
+    key_.Clear();
     restart_index_ = index;
     // current_ will be fixed by ParseNextKey();
 
@@ -143,7 +144,7 @@ class Block::Iter : public Iterator {
   virtual Status status() const { return status_; }
   virtual Slice key() const {
     assert(Valid());
-    return key_;
+    return key_.GetKey();
   }
   virtual Slice value() const {
     assert(Valid());
@@ -193,7 +194,7 @@ class Block::Iter : public Iterator {
     // Linear search (within restart block) for first key >= target
 
     while (true) {
-      if (!ParseNextKey() || Compare(key_, target) >= 0) {
+      if (!ParseNextKey() || Compare(key_.GetKey(), target) >= 0) {
         return;
       }
     }
@@ -215,7 +216,7 @@ class Block::Iter : public Iterator {
     current_ = restarts_;
     restart_index_ = num_restarts_;
     status_ = Status::Corruption("bad entry in block");
-    key_.clear();
+    key_.Clear();
     value_.clear();
   }
 
@@ -233,12 +234,11 @@ class Block::Iter : public Iterator {
     // Decode next entry
     uint32_t shared, non_shared, value_length;
     p = DecodeEntry(p, limit, &shared, &non_shared, &value_length);
-    if (p == nullptr || key_.size() < shared) {
+    if (p == nullptr || key_.Size() < shared) {
       CorruptionError();
       return false;
     } else {
-      key_.resize(shared);
+      key_.TrimAppend(shared, p, non_shared);
-      key_.append(p, non_shared);
       value_ = Slice(p + non_shared, value_length);
       while (restart_index_ + 1 < num_restarts_ &&
              GetRestartPoint(restart_index_ + 1) < current_) {