From 9d390bc5db4f8a15a1b91b0b9ffdd4cb9fd615cb Mon Sep 17 00:00:00 2001
From: Dev Ojha <ValarDragon@users.noreply.github.com>
Date: Tue, 6 Feb 2024 02:20:18 -0500
Subject: [PATCH] perf: CacheCtx speedups (#504)

* CacheKV speedups

* Remove extra btree alloc
---
 store/cachekv/search_benchmark_test.go |  7 +++----
 store/cachekv/store.go                 | 16 ++++++++++++----
 store/cachemulti/store.go              |  2 +-
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/store/cachekv/search_benchmark_test.go b/store/cachekv/search_benchmark_test.go
index 4007c7cda202..ddf4bdf02c15 100644
--- a/store/cachekv/search_benchmark_test.go
+++ b/store/cachekv/search_benchmark_test.go
@@ -3,8 +3,6 @@ package cachekv
 import (
 	"strconv"
 	"testing"
-
-	"github.com/cosmos/cosmos-sdk/store/cachekv/internal"
 )
 
 func BenchmarkLargeUnsortedMisses(b *testing.B) {
@@ -36,9 +34,10 @@ func generateStore() *Store {
 		cache[key] = &cValue{}
 	}
 
-	return &Store{
+	store := &Store{
 		cache:         cache,
 		unsortedCache: unsorted,
-		sortedCache:   internal.NewBTree(),
 	}
+	store.resetSortedCache()
+	return store
 }
diff --git a/store/cachekv/store.go b/store/cachekv/store.go
index a8f468979cd9..2cb5a63fa6a4 100644
--- a/store/cachekv/store.go
+++ b/store/cachekv/store.go
@@ -28,7 +28,7 @@ type Store struct {
 	mtx           sync.Mutex
 	cache         map[string]*cValue
 	unsortedCache map[string]struct{}
-	sortedCache   internal.BTree // always ascending sorted
+	sortedCache   *internal.BTree // always ascending sorted
 	parent        types.KVStore
 }
 
@@ -39,7 +39,7 @@ func NewStore(parent types.KVStore) *Store {
 	return &Store{
 		cache:         make(map[string]*cValue),
 		unsortedCache: make(map[string]struct{}),
-		sortedCache:   internal.NewBTree(),
+		sortedCache:   nil,
 		parent:        parent,
 	}
 }
@@ -93,13 +93,18 @@ func (store *Store) Delete(key []byte) {
 	store.setCacheValue(key, nil, true)
 }
 
+func (store *Store) resetSortedCache() {
+	newTree := internal.NewBTree()
+	store.sortedCache = &newTree
+}
+
 // Implements Cachetypes.KVStore.
 func (store *Store) Write() {
 	store.mtx.Lock()
 	defer store.mtx.Unlock()
 
 	if len(store.cache) == 0 && len(store.unsortedCache) == 0 {
-		store.sortedCache = internal.NewBTree()
+		store.sortedCache = nil
 		return
 	}
 
@@ -140,7 +145,7 @@ func (store *Store) Write() {
 	for key := range store.unsortedCache {
 		delete(store.unsortedCache, key)
 	}
-	store.sortedCache = internal.NewBTree()
+	store.sortedCache = nil
 }
 
 // CacheWrap implements CacheWrapper.
@@ -169,6 +174,9 @@ func (store *Store) ReverseIterator(start, end []byte) types.Iterator {
 func (store *Store) iterator(start, end []byte, ascending bool) types.Iterator {
 	store.mtx.Lock()
 	defer store.mtx.Unlock()
+	if store.sortedCache == nil {
+		store.resetSortedCache()
+	}
 
 	store.dirtyItems(start, end)
 	isoSortedCache := store.sortedCache.Copy()
diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
index 86927466a905..a162910fe5a5 100644
--- a/store/cachemulti/store.go
+++ b/store/cachemulti/store.go
@@ -73,7 +73,7 @@ func NewStore(
 }
 
 func newCacheMultiStoreFromCMS(cms Store) Store {
-	stores := make(map[types.StoreKey]types.CacheWrapper)
+	stores := make(map[types.StoreKey]types.CacheWrapper, len(cms.stores))
 	for k, v := range cms.stores {
 		stores[k] = v
 	}