Improve DocDB and IntentAwareIterator performance

Summary:
Improve DocDB and IntentAwareIterator performance by:
- avoids a few duplicate DocHybridTime decoding and KeyBytes encoding in DocRowwiseIterator::HasNext() and BuildSubDocument()
- avoids extra IntentAwareIterator::FetchKey() in DocRowwiseIterator::HasNext() to check for infinite loop.
- avoids extra buffer malloc in IntentAwareIterator::SeekOutOfSubDoc()
- reserves buffer IntentAwareIterator::SeekForward()

These changes bring about 15% improvement in CassandraPersonalization read workload.

Test Plan:
Before:

```
java -jar ~/code/yugabyte/java/yb-loadtester/target/yb-sample-apps.jar -workload CassandraPersonalization -nodes 127.0.0.1:9042 -nouuid -max_written_key 10 -read_only -num_threads_write 0 -num_threads_read 1
...
2018-04-17 16:51:56,662 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 38.20 ops/sec (26.21 ms/op), 1336 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 35021 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
2018-04-17 16:52:01,662 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 38.19 ops/sec (26.14 ms/op), 1527 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 40022 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
2018-04-17 16:52:06,662 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 37.80 ops/sec (26.45 ms/op), 1716 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 45022 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
```

After:

```
Uptime: 30020 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
2018-04-17 16:53:47,584 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 44.16 ops/sec (22.63 ms/op), 1548 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 35025 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
2018-04-17 16:53:52,589 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 44.55 ops/sec (22.51 ms/op), 1771 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 40030 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
2018-04-17 16:53:57,589 [INFO|com.yugabyte.sample.common.metrics.MetricsTracker|MetricsTracker] Read: 44.20 ops/sec (22.53 ms/op), 1992 total ops  |  Write: 0.00 ops/sec (0.00 ms/op), 0 total ops  |  Uptime: 45030 ms | maxWrittenKey: 10 | maxGeneratedKey: 10 |
```

Reviewers: sergei, mikhail

Reviewed By: mikhail

Subscribers: ybase

Differential Revision: https://phabricator.dev.yugabyte.com/D4539

src/yb/docdb/doc_key.cc

@@ -219,6 +219,12 @@ yb::Status DocKey::DecodeFrom(rocksdb::Slice *slice, DocKeyPart part_to_decode)
   return DoDecode(slice, part_to_decode, DecodeFromCallback(this));
 }
 
+Result<size_t> DocKey::DecodeFrom(const rocksdb::Slice& slice, DocKeyPart part_to_decode) {
+  rocksdb::Slice copy = slice;
+  RETURN_NOT_OK(DecodeFrom(&copy, part_to_decode));
+  return slice.size() - copy.size();
+}
+
 template<class Callback>
 yb::Status DocKey::DoDecode(rocksdb::Slice *slice,
                             DocKeyPart part_to_decode,

src/yb/docdb/doc_key.h

@@ -111,6 +111,11 @@ class DocKey {
   CHECKED_STATUS DecodeFrom(rocksdb::Slice* slice,
       DocKeyPart part_to_decode = DocKeyPart::WHOLE_DOC_KEY);
 
+  // Decodes a document key from the given RocksDB key similar to the above but return the number
+  // of bytes decoded from the input slice.
+  Result<size_t> DecodeFrom(const rocksdb::Slice& slice,
+      DocKeyPart part_to_decode = DocKeyPart::WHOLE_DOC_KEY);
+
   // Splits given RocksDB key into vector of slices that forms range_group of document key.
   static CHECKED_STATUS PartiallyDecode(rocksdb::Slice* slice,
                                         boost::container::small_vector_base<Slice>* out);

src/yb/docdb/doc_rowwise_iterator.cc

@@ -213,12 +213,25 @@ bool DocRowwiseIterator::HasNext() const {
       status_ = fetched_key.status();
       return true;
     }
-    {
-      Slice key_copy = *fetched_key;
-      status_ = row_key_.DecodeFrom(&key_copy);
+
+    // The iterator is positioned by the previous GetSubDocument call (which places the iterator
+    // outside the previous doc_key). Ensure the iterator is pushed forward/backward indeed. We
+    // check it here instead of after GetSubDocument() below because we want to avoid the extra
+    // expensive FetchKey() call just to fetch and validate the key.
+    if (!iter_key_.data().empty() &&
+        (is_forward_scan_ ? iter_key_.CompareTo(*fetched_key) >= 0
+                          : iter_key_.CompareTo(*fetched_key) <= 0)) {
+      status_ = STATUS_SUBSTITUTE(Corruption, "Infinite loop detected at $0",
+                                  FormatRocksDBSliceAsStr(*fetched_key));
+      VLOG(1) << status_;
+      return true;
     }
-    if (!status_.ok()) {
+    iter_key_.Reset(*fetched_key);
+
+    const Result<size_t> dockey_size = row_key_.DecodeFrom(iter_key_);
+    if (!dockey_size.ok()) {
       // Defer error reporting to NextRow().
+      status_ = dockey_size.status();
       return true;
     }
 
@@ -227,10 +240,8 @@ bool DocRowwiseIterator::HasNext() const {
       return false;
     }
 
-    KeyBytes old_key(*fetched_key);
-    // The iterator is positioned by the previous GetSubDocument call
-    // (which places the iterator outside the previous doc_key).
-    auto sub_doc_key = SubDocKey(row_key_).EncodeWithoutHt();
+    // Prepare the DocKey to get the SubDocument. Trim the DocKey to contain just the primary key.
+    Slice sub_doc_key(iter_key_.data().data(), *dockey_size);
     GetSubDocumentData data = { sub_doc_key, &row_, &doc_found };
     data.table_ttl = TableTTL(schema_);
     status_ = GetSubDocument(db_iter_.get(), data, &projection_subkeys_);
@@ -253,20 +264,6 @@ bool DocRowwiseIterator::HasNext() const {
         return true;
       }
     }
-    // GetSubDocument must ensure that iterator is pushed forward, to avoid loops.
-    if (db_iter_->valid()) {
-      auto iter_key = db_iter_->FetchKey();
-      if (!iter_key.ok()) {
-        status_ = iter_key.status();
-        return true;
-      }
-      if (old_key.AsSlice().compare(*iter_key) >= 0) {
-        status_ = STATUS_SUBSTITUTE(Corruption, "Infinite loop detected at $0",
-            FormatRocksDBSliceAsStr(old_key.AsSlice()));
-        VLOG(1) << status_;
-        return true;
-      }
-    }
     status_ = EnsureIteratorPositionCorrect();
     if (!status_.ok()) {
       // Defer error reporting to NextRow().

src/yb/docdb/doc_rowwise_iterator.h

@@ -170,12 +170,15 @@ class DocRowwiseIterator : public common::YQLRowwiseIteratorIf {
   // Indicates whether we've already finished iterating.
   mutable bool done_;
 
-  // HasNext constructs the whole row SubDocument.
+  // HasNext constructs the whole row's SubDocument.
   mutable SubDocument row_;
 
-  // The current row's Primary key. It is set to lower bound in the beginning.
+  // The current row's primary key. It is set to lower bound in the beginning.
   mutable DocKey row_key_;
 
+  // The current row's iterator key.
+  mutable KeyBytes iter_key_;
+
   // When HasNext constructs a row, row_ready_ is set to true.
   // When NextRow consumes the row, this variable is set to false.
   // It is initialized to false, to make sure first HasNext constructs a new row.

src/yb/docdb/docdb.cc

@@ -409,19 +409,14 @@ CHECKED_STATUS BuildSubDocument(
   while (iter->valid()) {
     // Since we modify num_values_observed on recursive calls, we keep a local copy of the value.
     int64 current_values_observed = *num_values_observed;
-    auto key = VERIFY_RESULT(iter->FetchKey());
+    DocHybridTime doc_ht;
+    auto key = VERIFY_RESULT(iter->FetchKey(&doc_ht));
     VLOG(4) << "iter: " << SubDocKey::DebugSliceToString(key)
             << ", key: " << SubDocKey::DebugSliceToString(data.subdocument_key);
     DCHECK(key.starts_with(data.subdocument_key))
         << "iter: " << SubDocKey::DebugSliceToString(key)
         << ", key: " << SubDocKey::DebugSliceToString(data.subdocument_key);
 
-    auto doc_ht = VERIFY_RESULT(DocHybridTime::DecodeFromEnd(&key));
-    if (key.empty() || static_cast<ValueType>(key[key.size() - 1]) != ValueType::kHybridTime) {
-      return STATUS_FORMAT(Corruption, "Key missing value type for hybrid time: $0",
-                           key.ToDebugHexString());
-    }
-    key.remove_suffix(1);
     // Key could be invalidated because we could move iterator, so back it up.
     KeyBytes key_copy(key);
     key = key_copy.AsSlice();
@@ -514,15 +509,15 @@ CHECKED_STATUS BuildSubDocument(
             ? write_time.hybrid_time().GetPhysicalValueMicros()
             : doc_value.user_timestamp());
         if (!data.high_index->CanInclude(current_values_observed)) {
-          iter->SeekOutOfSubDoc(key);
+          iter->SeekOutOfSubDoc(&key_copy);
           return Status::OK();
         }
         if (data.low_index->CanInclude(*num_values_observed)) {
           *data.result = SubDocument(doc_value.primitive_value());
         }
         (*num_values_observed)++;
         VLOG(3) << "SeekOutOfSubDoc: " << SubDocKey::DebugSliceToString(key);
-        iter->SeekOutOfSubDoc(key);
+        iter->SeekOutOfSubDoc(&key_copy);
         return Status::OK();
       }
     }

src/yb/docdb/docdb_rocksdb_util.cc

@@ -203,8 +203,16 @@ void SeekPastSubKey(const Slice& key, rocksdb::Iterator* iter) {
 }
 
 void SeekOutOfSubKey(const Slice& key, rocksdb::Iterator* iter) {
-  KeyBytes key_bytes(key, ValueTypeAsChar::kMaxByte);
-  SeekForward(key_bytes, iter);
+  KeyBytes key_bytes;
+  key_bytes.Reserve(key.size() + 1);
+  key_bytes.AppendRawBytes(key);
+  SeekOutOfSubKey(&key_bytes, iter);
+}
+
+void SeekOutOfSubKey(KeyBytes* key_bytes, rocksdb::Iterator* iter) {
+  key_bytes->AppendValueType(ValueType::kMaxByte);
+  SeekForward(*key_bytes, iter);
+  key_bytes->RemoveValueTypeSuffix(ValueType::kMaxByte);
 }
 
 void PerformRocksDBSeek(

src/yb/docdb/docdb_rocksdb_util.h

@@ -61,7 +61,11 @@ void SeekForward(const KeyBytes& key_bytes, rocksdb::Iterator *iter);
 void SeekPastSubKey(const SubDocKey& sub_doc_key, rocksdb::Iterator* iter);
 void SeekPastSubKey(const Slice& key, rocksdb::Iterator* iter);
 
+// Seek out of the given SubDocKey. For efficiency, the method that takes a non-const KeyBytes
+// pointer avoids memory allocation by using the KeyBytes buffer to prepare the key to seek to by
+// appending an extra byte. The appended byte is removed when the method returns.
 void SeekOutOfSubKey(const Slice& key, rocksdb::Iterator* iter);
+void SeekOutOfSubKey(KeyBytes* key_bytes, rocksdb::Iterator* iter);
 
 KeyBytes AppendDocHt(const Slice& key, const DocHybridTime& doc_ht);
 

src/yb/docdb/docrowwiseiterator-test.cc

@@ -941,12 +941,13 @@ SubDocKey(DocKey([], ["row2", 22222]), [ColumnId(40); HT{ physical: 1000 }]) ->
       rocksdb(), rocksdb::ReadOptions(), ReadHybridTime::FromMicros(1000), boost::none);
   iter.Seek(DocKey());
   ASSERT_TRUE(iter.valid());
-  Result<Slice> key = iter.FetchKey();
+  DocHybridTime doc_ht;
+  Result<Slice> key = iter.FetchKey(&doc_ht);
   ASSERT_OK(key);
   SubDocKey subdoc_key;
-  ASSERT_OK(subdoc_key.FullyDecodeFrom(*key));
-  ASSERT_EQ(subdoc_key.ToString(),
-            R"#(SubDocKey(DocKey([], ["row1", 11111]), [ColumnId(30); HT{ physical: 1000 }]))#");
+  ASSERT_OK(subdoc_key.FullyDecodeFrom(*key, HybridTimeRequired::kFalse));
+  ASSERT_EQ(subdoc_key.ToString(), R"#(SubDocKey(DocKey([], ["row1", 11111]), [ColumnId(30)]))#");
+  ASSERT_EQ(doc_ht.ToString(), "HT{ physical: 1000 }");
 }
 
 }  // namespace docdb

src/yb/docdb/intent_aware_iterator.cc

@@ -96,7 +96,7 @@ Result<HybridTime> TransactionStatusCache::DoGetCommitTime(const TransactionId&
     };
     txn_status_manager_->RequestStatusAt(
         {&transaction_id, read_time_.read, read_time_.global_limit, read_time_.serial_no,
-         callback});
+              callback});
     future.wait();
     auto txn_status_result = future.get();
     if (txn_status_result.ok()) {
@@ -125,7 +125,7 @@ Result<HybridTime> TransactionStatusCache::DoGetCommitTime(const TransactionId&
     return local_commit_time.is_valid() ? local_commit_time : HybridTime::kMin;
   } else {
     return txn_status.status == TransactionStatus::COMMITTED ? txn_status.status_time
-                                                             : HybridTime::kMin;
+        : HybridTime::kMin;
   }
 }
 
@@ -163,7 +163,7 @@ Result<DecodeStrongWriteIntentResult> DecodeStrongWriteIntent(
   DocHybridTime intent_ht;
   DecodeStrongWriteIntentResult result;
   RETURN_NOT_OK(DecodeIntentKey(
-     intent_iter->key(), &result.intent_prefix, &result.intent_type, &intent_ht));
+      intent_iter->key(), &result.intent_prefix, &result.intent_type, &intent_ht));
   if (IsStrongWriteIntent(result.intent_type)) {
     result.intent_value = intent_iter->value();
     auto txn_id = VERIFY_RESULT(DecodeTransactionIdFromIntentValue(&result.intent_value));
@@ -269,7 +269,11 @@ void IntentAwareIterator::Seek(const Slice& key) {
 }
 
 void IntentAwareIterator::SeekForward(const Slice& key) {
-  KeyBytes key_bytes(key);
+  KeyBytes key_bytes;
+  // Reserve space for key plus kMaxBytesPerEncodedHybridTime + 1 bytes for SeekForward() below to
+  // avoid extra realloc while appending the read time.
+  key_bytes.Reserve(key.size() + kMaxBytesPerEncodedHybridTime + 1);
+  key_bytes.AppendRawBytes(key);
   SeekForward(&key_bytes);
 }
 
@@ -316,25 +320,34 @@ void IntentAwareIterator::SeekPastSubKey(const Slice& key) {
   }
 }
 
-void IntentAwareIterator::SeekOutOfSubDoc(const Slice& key) {
-  VLOG(4) << "SeekOutOfSubDoc(" << SubDocKey::DebugSliceToString(key) << ")";
+void IntentAwareIterator::SeekOutOfSubDoc(KeyBytes* key_bytes) {
+  VLOG(4) << "SeekOutOfSubDoc(" << SubDocKey::DebugSliceToString(*key_bytes) << ")";
   if (!status_.ok()) {
     return;
   }
 
-  docdb::SeekOutOfSubKey(key, iter_.get());
+  docdb::SeekOutOfSubKey(key_bytes, iter_.get());
   if (intent_iter_ && status_.ok()) {
     status_ = SetIntentUpperbound();
     if (!status_.ok()) {
       return;
     }
-    KeyBytes intent_prefix = GetIntentPrefixForKeyWithoutHt(key);
+    GetIntentPrefixForKeyWithoutHt(*key_bytes, &seek_key_buffer_);
     // See comment for SubDocKey::AdvanceOutOfSubDoc.
-    intent_prefix.AppendValueType(ValueType::kMaxByte);
-    SeekForwardToSuitableIntent(intent_prefix);
+    seek_key_buffer_.AppendValueType(ValueType::kMaxByte);
+    SeekForwardToSuitableIntent(seek_key_buffer_);
   }
 }
 
+void IntentAwareIterator::SeekOutOfSubDoc(const Slice& key) {
+  KeyBytes key_bytes;
+  // Reserve space for key + 1 byte for docdb::SeekOutOfSubKey() above to avoid extra realloc while
+  // appending kMaxByte.
+  key_bytes.Reserve(key.size() + 1);
+  key_bytes.AppendRawBytes(key);
+  SeekOutOfSubDoc(&key_bytes);
+}
+
 void IntentAwareIterator::SeekToLastDocKey() {
   if (intent_iter_) {
     // TODO (dtxn): Implement SeekToLast when inten intents are present. Since part of the
@@ -401,20 +414,26 @@ bool IntentAwareIterator::IsEntryRegular() {
   return true;
 }
 
-Result<Slice> IntentAwareIterator::FetchKey() {
+Result<Slice> IntentAwareIterator::FetchKey(DocHybridTime* doc_ht) {
   RETURN_NOT_OK(status_);
   Slice result;
+  DocHybridTime doc_ht_seen;
   if (IsEntryRegular()) {
     result = iter_->key();
+    doc_ht_seen = VERIFY_RESULT(DocHybridTime::DecodeFromEnd(&result));
+    DCHECK(result.ends_with(ValueTypeAsChar::kHybridTime)) << result.ToDebugString();
+    result.remove_suffix(1);
   } else {
     DCHECK_EQ(ResolvedIntentState::kValid, resolved_intent_state_);
-    result = resolved_intent_sub_doc_key_encoded_;
+    result = resolved_intent_sub_doc_key_;
+    doc_ht_seen = resolved_intent_txn_dht_;
   }
-  DocHybridTime doc_ht;
-  RETURN_NOT_OK(DecodeHybridTimeFromEndOfKey(result, &doc_ht));
-  max_seen_ht_.MakeAtLeast(doc_ht.hybrid_time());
+  if (doc_ht != nullptr) {
+    *doc_ht = doc_ht_seen;
+  }
+  max_seen_ht_.MakeAtLeast(doc_ht_seen.hybrid_time());
   VLOG(4) << "Fetched key " << SubDocKey::DebugSliceToString(result)
-          <<  ", with time: " << doc_ht.hybrid_time()
+          << ", with time: " << doc_ht_seen
           << ", while read bounds are: " << read_time_;
   return result;
 }
@@ -452,7 +471,7 @@ void IntentAwareIterator::ProcessIntent() {
                                                    : resolved_intent_txn_dht_;
   if (decode_result->value_time > real_time &&
       (decode_result->same_transaction ||
-           decode_result->value_time.hybrid_time() <= read_time_.global_limit)) {
+       decode_result->value_time.hybrid_time() <= read_time_.global_limit)) {
     if (resolved_intent_state_ == ResolvedIntentState::kNoIntent) {
       resolved_intent_key_prefix_.Reset(decode_result->intent_prefix);
       auto prefix = prefix_stack_.empty() ? Slice() : prefix_stack_.back();
@@ -463,7 +482,7 @@ void IntentAwareIterator::ProcessIntent() {
       }
       resolved_intent_state_ =
           decode_result->intent_prefix.starts_with(prefix) ? ResolvedIntentState::kValid
-                                                           : ResolvedIntentState::kInvalidPrefix;
+          : ResolvedIntentState::kInvalidPrefix;
     }
     if (decode_result->same_transaction) {
       intent_dht_from_same_txn_ = decode_result->value_time;
@@ -476,8 +495,12 @@ void IntentAwareIterator::ProcessIntent() {
 }
 
 void IntentAwareIterator::UpdateResolvedIntentSubDocKeyEncoded() {
-  resolved_intent_sub_doc_key_encoded_.ResetRawBytes(
-      resolved_intent_key_prefix_.data().data() + 1, resolved_intent_key_prefix_.size() - 1);
+  resolved_intent_sub_doc_key_ = resolved_intent_key_prefix_;
+  status_ = resolved_intent_sub_doc_key_.consume_byte(ValueTypeAsChar::kIntentPrefix);
+  if (!status_.ok()) {
+    return;
+  }
+  resolved_intent_sub_doc_key_encoded_.Reset(resolved_intent_sub_doc_key_);
   resolved_intent_sub_doc_key_encoded_.AppendValueType(ValueType::kHybridTime);
   resolved_intent_sub_doc_key_encoded_.AppendHybridTime(resolved_intent_txn_dht_);
   VLOG(4) << "Resolved intent SubDocKey: "
@@ -551,9 +574,10 @@ void IntentAwareIterator::DebugDump() {
   }
   LOG(INFO) << "valid(): " << valid();
   if (valid()) {
-    auto key = FetchKey();
+    DocHybridTime doc_ht;
+    auto key = FetchKey(&doc_ht);
     if (key.ok()) {
-      LOG(INFO) << "key(): " << DebugDumpKeyToStr(*key);
+      LOG(INFO) << "key(): " << DebugDumpKeyToStr(*key) << ", doc_ht: " << doc_ht;
     } else {
       LOG(INFO) << "key(): fetch failed: " << key.status();
     }
@@ -571,7 +595,7 @@ Status IntentAwareIterator::FindLastWriteTime(
 
   DOCDB_DEBUG_SCOPE_LOG(
       SubDocKey::DebugSliceToString(key_without_ht) + ", " + yb::ToString(max_deleted_ts) + ", "
-          + yb::ToString(result_value),
+      + yb::ToString(result_value),
       std::bind(&IntentAwareIterator::DebugDump, this));
   DCHECK(!DebugHasHybridTime(key_without_ht));
 
@@ -649,7 +673,7 @@ void IntentAwareIterator::PopPrefix() {
   skip_future_intents_needed_ = true;
   VLOG(4) << "PopPrefix: "
           << (prefix_stack_.empty() ? std::string()
-                                    : SubDocKey::DebugSliceToString(prefix_stack_.back()));
+              : SubDocKey::DebugSliceToString(prefix_stack_.back()));
 }
 
 void IntentAwareIterator::SkipFutureRecords() {
@@ -701,16 +725,10 @@ void IntentAwareIterator::SkipFutureIntents() {
   }
   auto prefix = prefix_stack_.empty() ? Slice() : prefix_stack_.back();
   if (resolved_intent_state_ != ResolvedIntentState::kNoIntent) {
-    VLOG(4) << "Checking resolved intent: "
-            << resolved_intent_key_prefix_.AsSlice().ToDebugHexString()
+    VLOG(4) << "Checking resolved intent subdockey: "
+            << resolved_intent_sub_doc_key_.ToDebugHexString()
             << ", against new prefix: " << prefix.ToDebugHexString();
-    auto resolved_intent_key_prefix = resolved_intent_key_prefix_.AsSlice();
-    status_ = resolved_intent_key_prefix.consume_byte(ValueTypeAsChar::kIntentPrefix);
-    if (!status_.ok()) {
-      status_ = status_.CloneAndPrepend("Bad resolved intent key");
-      return;
-    }
-    auto compare_result = resolved_intent_key_prefix.compare_prefix(prefix);
+    auto compare_result = resolved_intent_sub_doc_key_.compare_prefix(prefix);
     if (compare_result == 0) {
       resolved_intent_state_ = ResolvedIntentState::kValid;
       return;