perf: Speedup cachekv iterator on large deletions & IBC v2 upgrade lo…

…gic (backport #10741) (#10745)

* perf: Speedup cachekv iterator on large deletions & IBC v2 upgrade logic (#10741)

(cherry picked from commit 314e1d5)

# Conflicts:
#	CHANGELOG.md
#	store/cachekv/store_bench_test.go

* fix conflicts

Co-authored-by: Dev Ojha <[email protected]>
Co-authored-by: marbar3778 <[email protected]>

CHANGELOG.md

@@ -40,7 +40,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ## [v0.44.5](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.44.5) - 2021-12-02
 
 ### Improvements
+
 * (baseapp) [\#10631](https://github.com/cosmos/cosmos-sdk/pull/10631)  Emit ante events even for the failed txs. 
+* (store) [\#10741](https://github.com/cosmos/cosmos-sdk/pull/10741) Significantly speedup iterator creation after delete heavy workloads. Significantly improves IBC migration times.
 
 ### Bug Fixes
 

store/cachekv/bench_helper_test.go

@@ -0,0 +1,44 @@
+package cachekv_test
+
+import "crypto/rand"
+
+func randSlice(sliceSize int) []byte {
+	bz := make([]byte, sliceSize)
+	_, _ = rand.Read(bz)
+	return bz
+}
+
+func incrementByteSlice(bz []byte) {
+	for index := len(bz) - 1; index >= 0; index-- {
+		if bz[index] < 255 {
+			bz[index]++
+			break
+		} else {
+			bz[index] = 0
+		}
+	}
+}
+
+// Generate many keys starting at startKey, and are in sequential order
+func generateSequentialKeys(startKey []byte, numKeys int) [][]byte {
+	toReturn := make([][]byte, 0, numKeys)
+	cur := make([]byte, len(startKey))
+	copy(cur, startKey)
+	for i := 0; i < numKeys; i++ {
+		newKey := make([]byte, len(startKey))
+		copy(newKey, cur)
+		toReturn = append(toReturn, newKey)
+		incrementByteSlice(cur)
+	}
+	return toReturn
+}
+
+// Generate many random, unsorted keys
+func generateRandomKeys(keySize int, numKeys int) [][]byte {
+	toReturn := make([][]byte, 0, numKeys)
+	for i := 0; i < numKeys; i++ {
+		newKey := randSlice(keySize)
+		toReturn = append(toReturn, newKey)
+	}
+	return toReturn
+}

store/cachekv/memiterator.go

@@ -1,6 +1,8 @@
 package cachekv
 
 import (
+	"bytes"
+
 	dbm "github.com/tendermint/tm-db"
 
 	"github.com/cosmos/cosmos-sdk/store/types"
@@ -12,6 +14,7 @@ import (
 type memIterator struct {
 	types.Iterator
 
+	lastKey []byte
 	deleted map[string]struct{}
 }
 
@@ -29,22 +32,25 @@ func newMemIterator(start, end []byte, items *dbm.MemDB, deleted map[string]stru
 		panic(err)
 	}
 
-	newDeleted := make(map[string]struct{})
-	for k, v := range deleted {
-		newDeleted[k] = v
-	}
-
 	return &memIterator{
 		Iterator: iter,
 
-		deleted: newDeleted,
+		lastKey: nil,
+		deleted: deleted,
 	}
 }
 
 func (mi *memIterator) Value() []byte {
 	key := mi.Iterator.Key()
-	if _, ok := mi.deleted[string(key)]; ok {
+	// We need to handle the case where deleted is modified and includes our current key
+	// We handle this by maintaining a lastKey object in the iterator.
+	// If the current key is the same as the last key (and last key is not nil / the start)
+	// then we are calling value on the same thing as last time.
+	// Therefore we don't check the mi.deleted to see if this key is included in there.
+	reCallingOnOldLastKey := (mi.lastKey != nil) && bytes.Equal(key, mi.lastKey)
+	if _, ok := mi.deleted[string(key)]; ok && !reCallingOnOldLastKey {
 		return nil
 	}
+	mi.lastKey = key
 	return mi.Iterator.Value()
 }

store/cachekv/store_bench_test.go

@@ -1,8 +1,6 @@
 package cachekv_test
 
 import (
-	"crypto/rand"
-	"sort"
 	"testing"
 
 	dbm "github.com/tendermint/tm-db"
@@ -11,37 +9,136 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/dbadapter"
 )
 
-func benchmarkCacheKVStoreIterator(numKVs int, b *testing.B) {
+var sink interface{}
+
+const defaultValueSizeBz = 1 << 12
+
+// This benchmark measures the time of iterator.Next() when the parent store is blank
+func benchmarkBlankParentIteratorNext(b *testing.B, keysize int) {
+	mem := dbadapter.Store{DB: dbm.NewMemDB()}
+	kvstore := cachekv.NewStore(mem)
+	// Use a singleton for value, to not waste time computing it
+	value := randSlice(defaultValueSizeBz)
+	// Use simple values for keys, pick a random start,
+	// and take next b.N keys sequentially after.]
+	startKey := randSlice(32)
+
+	// Add 1 to avoid issues when b.N = 1
+	keys := generateSequentialKeys(startKey, b.N+1)
+	for _, k := range keys {
+		kvstore.Set(k, value)
+	}
+
 	b.ReportAllocs()
+	b.ResetTimer()
+
+	iter := kvstore.Iterator(keys[0], keys[b.N])
+	defer iter.Close()
+
+	for _ = iter.Key(); iter.Valid(); iter.Next() {
+		// deadcode elimination stub
+		sink = iter
+	}
+}
+
+// Benchmark setting New keys to a store, where the new keys are in sequence.
+func benchmarkBlankParentAppend(b *testing.B, keysize int) {
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
-	cstore := cachekv.NewStore(mem)
-	keys := make([]string, numKVs)
+	kvstore := cachekv.NewStore(mem)
+
+	// Use a singleton for value, to not waste time computing it
+	value := randSlice(32)
+	// Use simple values for keys, pick a random start,
+	// and take next b.N keys sequentially after.
+	startKey := randSlice(32)
 
-	for i := 0; i < numKVs; i++ {
-		key := make([]byte, 32)
-		value := make([]byte, 32)
+	keys := generateSequentialKeys(startKey, b.N)
 
-		_, _ = rand.Read(key)
-		_, _ = rand.Read(value)
+	b.ReportAllocs()
+	b.ResetTimer()
 
-		keys[i] = string(key)
-		cstore.Set(key, value)
+	for _, k := range keys {
+		kvstore.Set(k, value)
 	}
+}
 
-	sort.Strings(keys)
+// Benchmark setting New keys to a store, where the new keys are random.
+// the speed of this function does not depend on the values in the parent store
+func benchmarkRandomSet(b *testing.B, keysize int) {
+	mem := dbadapter.Store{DB: dbm.NewMemDB()}
+	kvstore := cachekv.NewStore(mem)
 
-	for n := 0; n < b.N; n++ {
-		iter := cstore.Iterator([]byte(keys[0]), []byte(keys[numKVs-1]))
+	// Use a singleton for value, to not waste time computing it
+	value := randSlice(defaultValueSizeBz)
+	keys := generateRandomKeys(keysize, b.N)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for _, k := range keys {
+		kvstore.Set(k, value)
+	}
 
-		for _ = iter.Key(); iter.Valid(); iter.Next() {
-		}
+	iter := kvstore.Iterator(keys[0], keys[b.N])
+	defer iter.Close()
 
-		iter.Close()
+	for _ = iter.Key(); iter.Valid(); iter.Next() {
+		// deadcode elimination stub
+		sink = iter
 	}
 }
 
-func BenchmarkCacheKVStoreIterator500(b *testing.B)    { benchmarkCacheKVStoreIterator(500, b) }
-func BenchmarkCacheKVStoreIterator1000(b *testing.B)   { benchmarkCacheKVStoreIterator(1000, b) }
-func BenchmarkCacheKVStoreIterator10000(b *testing.B)  { benchmarkCacheKVStoreIterator(10000, b) }
-func BenchmarkCacheKVStoreIterator50000(b *testing.B)  { benchmarkCacheKVStoreIterator(50000, b) }
-func BenchmarkCacheKVStoreIterator100000(b *testing.B) { benchmarkCacheKVStoreIterator(100000, b) }
+// Benchmark creating an iterator on a parent with D entries,
+// that are all deleted in the cacheKV store.
+// We essentially are benchmarking the cacheKV iterator creation & iteration times
+// with the number of entries deleted in the parent.
+func benchmarkIteratorOnParentWithManyDeletes(b *testing.B, numDeletes int) {
+	mem := dbadapter.Store{DB: dbm.NewMemDB()}
+
+	// Use a singleton for value, to not waste time computing it
+	value := randSlice(32)
+	// Use simple values for keys, pick a random start,
+	// and take next D keys sequentially after.
+	startKey := randSlice(32)
+	keys := generateSequentialKeys(startKey, numDeletes)
+	// setup parent db with D keys.
+	for _, k := range keys {
+		mem.Set(k, value)
+	}
+	kvstore := cachekv.NewStore(mem)
+	// Delete all keys from the cache KV store.
+	// The keys[1:] is to keep at least one entry in parent, due to a bug in the SDK iterator design.
+	// Essentially the iterator will never be valid, in that it should never run.
+	// However, this is incompatible with the for loop structure the SDK uses, hence
+	// causes a panic. Thus we do keys[1:].
+	for _, k := range keys[1:] {
+		kvstore.Delete(k)
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	iter := kvstore.Iterator(keys[0], keys[b.N])
+	defer iter.Close()
+
+	for _ = iter.Key(); iter.Valid(); iter.Next() {
+		// deadcode elimination stub
+		sink = iter
+	}
+}
+
+func BenchmarkBlankParentIteratorNextKeySize32(b *testing.B) {
+	benchmarkBlankParentIteratorNext(b, 32)
+}
+
+func BenchmarkBlankParentAppendKeySize32(b *testing.B) {
+	benchmarkBlankParentAppend(b, 32)
+}
+
+func BenchmarkSetKeySize32(b *testing.B) {
+	benchmarkRandomSet(b, 32)
+}
+
+func BenchmarkIteratorOnParentWith1MDeletes(b *testing.B) {
+	benchmarkIteratorOnParentWithManyDeletes(b, 1_000_000)
+}