From 70dad718081d51b0462d7ae872516ff616ae6890 Mon Sep 17 00:00:00 2001
From: Peter Mattis <petermattis@gmail.com>
Date: Thu, 2 Jun 2016 13:51:34 -0400
Subject: [PATCH] storage/engine: reenable the Seek optimization
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Primarily affects conditional-put heavy transactions where there is an
existing value.

Fixes #6980.

name                    old time/op  new time/op  delta
Update1_Cockroach-8      356µs ± 2%   348µs ± 1%  -2.42%  (p=0.000 n=10+10)
Update10_Cockroach-8     647µs ± 2%   619µs ± 2%  -4.33%   (p=0.000 n=9+10)
Update100_Cockroach-8   2.92ms ± 2%  2.67ms ± 3%  -8.54%  (p=0.000 n=10+10)
Update1000_Cockroach-8  24.5ms ± 1%  22.3ms ± 2%  -8.91%  (p=0.000 n=10+10)
---
 storage/engine/rocksdb.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/storage/engine/rocksdb.go b/storage/engine/rocksdb.go
index 6500223bc260..168cbb084ec0 100644
--- a/storage/engine/rocksdb.go
+++ b/storage/engine/rocksdb.go
@@ -672,12 +672,16 @@ func (r *rocksDBBatch) flushMutations() {
 	if err := r.ApplyBatchRepr(r.builder.Finish()); err != nil {
 		panic(err)
 	}
+	// Force a seek of the underlying iterator on the next Seek/ReverseSeek.
+	r.prefixIter.reseek = true
+	r.normalIter.reseek = true
 }
 
 type rocksDBIterator struct {
 	engine Reader
 	iter   *C.DBIterator
 	valid  bool
+	reseek bool
 	key    C.DBKey
 	value  C.DBSlice
 }
@@ -733,9 +737,10 @@ func (r *rocksDBIterator) Seek(key MVCCKey) {
 		// to access start[0] in an explicit seek.
 		r.setState(C.DBIterSeekToFirst(r.iter))
 	} else {
-		// It's tempting to avoid seeking if we're already at the desired key,
-		// but it may already have changed on the underlying engine, so we
-		// must seek again.
+		// We can avoid seeking if we're already at the key we seek.
+		if r.valid && !r.reseek && key.Equal(r.unsafeKey()) {
+			return
+		}
 		r.setState(C.DBIterSeek(r.iter, goToCKey(key)))
 	}
 }
@@ -745,6 +750,10 @@ func (r *rocksDBIterator) SeekReverse(key MVCCKey) {
 	if len(key.Key) == 0 {
 		r.setState(C.DBIterSeekToLast(r.iter))
 	} else {
+		// We can avoid seeking if we're already at the key we seek.
+		if r.valid && !r.reseek && key.Equal(r.unsafeKey()) {
+			return
+		}
 		r.setState(C.DBIterSeek(r.iter, goToCKey(key)))
 		// Maybe the key sorts after the last key in RocksDB.
 		if !r.Valid() {
@@ -820,6 +829,7 @@ func (r *rocksDBIterator) Less(key MVCCKey) bool {
 
 func (r *rocksDBIterator) setState(state C.DBIterState) {
 	r.valid = bool(state.valid)
+	r.reseek = false
 	r.key = state.key
 	r.value = state.value
 }