From 51dd21926c677ae4a63c8f45992903e7b30f0d13 Mon Sep 17 00:00:00 2001
From: Siying Dong <siying.d@fb.com>
Date: Tue, 14 Jan 2014 10:42:36 -0800
Subject: [PATCH] DB::Put() to estimate write batch data size needed and
 pre-allocate buffer

Summary:
In one of CPU profiles, we see some CPU costs of string::reserve() inside Batch.Put(). This patch should be able to reduce some of the costs by allocating sufficient buffer before hand.

Since it is a trivial percentage of CPU costs, I didn't find a way to show the improvement in one of the benchmarks. I'll deploy it to same application and do the same CPU profiling to make sure those CPU costs are reduced.

Test Plan: make all check

Reviewers: haobo, kailiu, igor

Reviewed By: haobo

CC: leveldb, nkg-

Differential Revision: https://reviews.facebook.net/D15135
---
 db/db_impl.cc                 | 5 ++++-
 db/write_batch.cc             | 3 ++-
 include/rocksdb/write_batch.h | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index e7f2abf99..12e07868f 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -3758,7 +3758,10 @@ Status DBImpl::GetDbIdentity(std::string& identity) {
 // Default implementations of convenience methods that subclasses of DB
 // can call if they wish
 Status DB::Put(const WriteOptions& opt, const Slice& key, const Slice& value) {
-  WriteBatch batch;
+  // Pre-allocate size of write batch conservatively.
+  // 8 bytes are taken by header, 4 bytes for count, 1 byte for type,
+  // and we allocate 11 extra bytes for key length, as well as value length.
+  WriteBatch batch(key.size() + value.size() + 24);
   batch.Put(key, value);
   return Write(opt, &batch);
 }
diff --git a/db/write_batch.cc b/db/write_batch.cc
index 2cfc8bd7d..7a6106afa 100644
--- a/db/write_batch.cc
+++ b/db/write_batch.cc
@@ -36,7 +36,8 @@ namespace rocksdb {
 // WriteBatch header has an 8-byte sequence number followed by a 4-byte count.
 static const size_t kHeader = 12;
 
-WriteBatch::WriteBatch() {
+WriteBatch::WriteBatch(size_t reserved_bytes) {
+  rep_.reserve((reserved_bytes > kHeader) ? reserved_bytes : kHeader);
   Clear();
 }
 
diff --git a/include/rocksdb/write_batch.h b/include/rocksdb/write_batch.h
index 30abead50..e7ce16005 100644
--- a/include/rocksdb/write_batch.h
+++ b/include/rocksdb/write_batch.h
@@ -35,7 +35,7 @@ struct SliceParts;
 
 class WriteBatch {
  public:
-  WriteBatch();
+  explicit WriteBatch(size_t reserved_bytes = 0);
   ~WriteBatch();
 
   // Store the mapping "key->value" in the database.