perf(store): replace sync.Map with plain map in cachekv stores

giga/deps/store/cachekv.go

@@ -13,8 +13,8 @@
 // Store wraps an in-memory cache around an underlying types.KVStore.
 type Store struct {
 	mtx       sync.RWMutex
-	cache     *sync.Map
-	deleted   *sync.Map
+	cache     map[string]*types.CValue
+	deleted   map[string]struct{}
 	parent    types.KVStore
 	storeKey  types.StoreKey
 	cacheSize int
@@ -25,8 +25,8 @@
 // NewStore creates a new Store object
 func NewStore(parent types.KVStore, storeKey types.StoreKey, cacheSize int) *Store {
 	return &Store{
-		cache:     &sync.Map{},
-		deleted:   &sync.Map{},
+		cache:     make(map[string]*types.CValue),
+		deleted:   make(map[string]struct{}),
 		parent:    parent,
 		storeKey:  storeKey,
 		cacheSize: cacheSize,
@@ -44,8 +44,8 @@
 
 // getFromCache queries the write-through cache for a value by key.
 func (store *Store) getFromCache(key []byte) []byte {
-	if cv, ok := store.cache.Load(UnsafeBytesToStr(key)); ok {
-		return cv.(*types.CValue).Value()
+	if cv, ok := store.cache[UnsafeBytesToStr(key)]; ok {
+		return cv.Value()
 	}
 	return store.parent.Get(key)
 }
@@ -84,12 +84,11 @@
 	// Not the best, but probably not a bottleneck depending.
 	keys := []string{}
 
-	store.cache.Range(func(key, value any) bool {
-		if value.(*types.CValue).Dirty() {
-			keys = append(keys, key.(string))
+	for key, value := range store.cache {
+		if value.Dirty() {
+			keys = append(keys, key)
 		}
-		return true
-	})
+	}
 	sort.Strings(keys)
 	// TODO: Consider allowing usage of Batch, which would allow the write to
 	// at least happen atomically.
@@ -103,10 +102,10 @@
 			continue
 		}
 
-		cacheValue, ok := store.cache.Load(key)
-		if ok && cacheValue.(*types.CValue).Value() != nil {
+		cacheValue, ok := store.cache[key]
+		if ok && cacheValue.Value() != nil {
 			// It already exists in the parent, hence delete it.
-			store.parent.Set([]byte(key), cacheValue.(*types.CValue).Value())
+			store.parent.Set([]byte(key), cacheValue.Value())
 		}
 	}
 
@@ -115,14 +114,12 @@
 	// writes immediately visible until Commit(). By keeping the cache populated
 	// with clean entries, subsequent reads will still hit the cache instead of
 	// falling through to the parent which can't read uncommitted data.
-	store.cache.Range(func(key, value any) bool {
-		cv := value.(*types.CValue)
+	for key, cv := range store.cache {
 		// Replace with a clean (non-dirty) version of the same value
-		store.cache.Store(key, types.NewCValue(cv.Value(), false))
-		return true
-	})
+		store.cache[key] = types.NewCValue(cv.Value(), false)
+	}
 	// Clear the deleted map since those deletes have been sent to parent
-	store.deleted = &sync.Map{}
+	store.deleted = make(map[string]struct{})
 }
 
 // CacheWrap implements CacheWrapper.
@@ -144,16 +141,16 @@
 	types.AssertValidKey(key)
 
 	keyStr := UnsafeBytesToStr(key)
-	store.cache.Store(keyStr, types.NewCValue(value, dirty))
+	store.cache[keyStr] = types.NewCValue(value, dirty)
 	if deleted {
-		store.deleted.Store(keyStr, struct{}{})
+		store.deleted[keyStr] = struct{}{}
 	} else {
-		store.deleted.Delete(keyStr)
+		delete(store.deleted, keyStr)
 	}
 }
 
 func (store *Store) isDeleted(key string) bool {
-	_, ok := store.deleted.Load(key)
+	_, ok := store.deleted[key]
 	return ok
 }
 
@@ -173,20 +170,18 @@
 	for _, pk := range store.parent.GetAllKeyStrsInRange(start, end) {
 		keyStrs[pk] = struct{}{}
 	}
-	store.cache.Range(func(key, value any) bool {
-		kbz := []byte(key.(string))
+	for key, cv := range store.cache {
+		kbz := []byte(key)
 		if bytes.Compare(kbz, start) < 0 || bytes.Compare(kbz, end) >= 0 {
 			// we don't want to break out of the iteration since cache isn't sorted
-			return true
+			continue
 		}
-		cv := value.(*types.CValue)
 		if cv.Value() == nil {
-			delete(keyStrs, key.(string))
+			delete(keyStrs, key)
 		} else {
-			keyStrs[key.(string)] = struct{}{}
+			keyStrs[key] = struct{}{}
 		}
-		return true
-	})
+	}
 	for k := range keyStrs {
 		res = append(res, k)
 	}

sei-cosmos/store/cachekv/memiterator.go

@@ -2,7 +2,6 @@ package cachekv
 
 import (
 	"bytes"
-	"sync"
 
 	dbm "github.com/tendermint/tm-db"
 
@@ -16,13 +15,13 @@ type memIterator struct {
 	types.Iterator
 
 	lastKey []byte
-	deleted *sync.Map
+	deleted map[string]struct{}
 }
 
 func newMemIterator(
 	start, end []byte,
 	items *dbm.MemDB,
-	deleted *sync.Map,
+	deleted map[string]struct{},
 	ascending bool,
 ) *memIterator {
 	var iter types.Iterator
@@ -56,7 +55,7 @@ func (mi *memIterator) Value() []byte {
 	// 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.Load(string(key)); ok && !reCallingOnOldLastKey {
+	if _, ok := mi.deleted[string(key)]; ok && !reCallingOnOldLastKey {
 		return nil
 	}
 	mi.lastKey = key

sei-cosmos/store/cachekv/store.go

@@ -16,9 +16,9 @@
 // Store wraps an in-memory cache around an underlying types.KVStore.
 type Store struct {
 	mtx           sync.RWMutex
-	cache         *sync.Map
-	deleted       *sync.Map
-	unsortedCache *sync.Map
+	cache         map[string]*types.CValue
+	deleted       map[string]struct{}
+	unsortedCache map[string]struct{}
 	sortedCache   *dbm.MemDB // always ascending sorted
 	parent        types.KVStore
 	storeKey      types.StoreKey
@@ -30,9 +30,9 @@
 // NewStore creates a new Store object
 func NewStore(parent types.KVStore, storeKey types.StoreKey, cacheSize int) *Store {
 	return &Store{
-		cache:         &sync.Map{},
-		deleted:       &sync.Map{},
-		unsortedCache: &sync.Map{},
+		cache:         make(map[string]*types.CValue),
+		deleted:       make(map[string]struct{}),
+		unsortedCache: make(map[string]struct{}),
 		sortedCache:   nil,
 		parent:        parent,
 		storeKey:      storeKey,
@@ -51,8 +51,8 @@
 
 // getFromCache queries the write-through cache for a value by key.
 func (store *Store) getFromCache(key []byte) []byte {
-	if cv, ok := store.cache.Load(conv.UnsafeBytesToStr(key)); ok {
-		return cv.(*types.CValue).Value()
+	if cv, ok := store.cache[conv.UnsafeBytesToStr(key)]; ok {
+		return cv.Value()
 	}
 	return store.parent.Get(key)
 }
@@ -91,12 +91,11 @@
 	// Not the best, but probably not a bottleneck depending.
 	keys := []string{}
 
-	store.cache.Range(func(key, value any) bool {
-		if value.(*types.CValue).Dirty() {
-			keys = append(keys, key.(string))
+	for key, value := range store.cache {
+		if value.Dirty() {
+			keys = append(keys, key)
 		}
-		return true
-	})
+	}
 	sort.Strings(keys)
 	// TODO: Consider allowing usage of Batch, which would allow the write to
 	// at least happen atomically.
@@ -110,16 +109,16 @@
 			continue
 		}
 
-		cacheValue, ok := store.cache.Load(key)
-		if ok && cacheValue.(*types.CValue).Value() != nil {
+		cacheValue, ok := store.cache[key]
+		if ok && cacheValue.Value() != nil {
 			// It already exists in the parent, hence delete it.
-			store.parent.Set([]byte(key), cacheValue.(*types.CValue).Value())
+			store.parent.Set([]byte(key), cacheValue.Value())
 		}
 	}
 
-	store.cache = &sync.Map{}
-	store.deleted = &sync.Map{}
-	store.unsortedCache = &sync.Map{}
+	store.cache = make(map[string]*types.CValue)
+	store.deleted = make(map[string]struct{})
+	store.unsortedCache = make(map[string]struct{})
 	store.sortedCache = nil
 }
 
@@ -281,16 +280,13 @@
 	// Even without that, too many range checks eventually becomes more expensive
 	// than just not having the cache.
 	// store.emitUnsortedCacheSizeMetric()
-	store.unsortedCache.Range(func(key, value any) bool {
-		cKey := key.(string)
+	for cKey := range store.unsortedCache {
 		if dbm.IsKeyInDomain(conv.UnsafeStrToBytes(cKey), start, end) {
-			cacheValue, ok := store.cache.Load(key)
-			if ok {
-				unsorted = append(unsorted, &kv.Pair{Key: []byte(cKey), Value: cacheValue.(*types.CValue).Value()})
+			if cacheValue, ok := store.cache[cKey]; ok {
+				unsorted = append(unsorted, &kv.Pair{Key: []byte(cKey), Value: cacheValue.Value()})
 			}
 		}
-		return true
-	})
+	}
 	store.clearUnsortedCacheSubset(unsorted, stateUnsorted)
 	return
 }
@@ -323,7 +319,7 @@
 func (store *Store) deleteKeysFromUnsortedCache(unsorted []*kv.Pair) {
 	for _, kv := range unsorted {
 		keyStr := conv.UnsafeBytesToStr(kv.Key)
-		store.unsortedCache.Delete(keyStr)
+		delete(store.unsortedCache, keyStr)
 	}
 }
 
@@ -335,19 +331,19 @@
 	types.AssertValidKey(key)
 
 	keyStr := conv.UnsafeBytesToStr(key)
-	store.cache.Store(keyStr, types.NewCValue(value, dirty))
+	store.cache[keyStr] = types.NewCValue(value, dirty)
 	if deleted {
-		store.deleted.Store(keyStr, struct{}{})
+		store.deleted[keyStr] = struct{}{}
 	} else {
-		store.deleted.Delete(keyStr)
+		delete(store.deleted, keyStr)
 	}
 	if dirty {
-		store.unsortedCache.Store(keyStr, struct{}{})
+		store.unsortedCache[keyStr] = struct{}{}
 	}
 }
 
 func (store *Store) isDeleted(key string) bool {
-	_, ok := store.deleted.Load(key)
+	_, ok := store.deleted[key]
 	return ok
 }
 
@@ -367,20 +363,18 @@
 	for _, pk := range store.parent.GetAllKeyStrsInRange(start, end) {
 		keyStrs[pk] = struct{}{}
 	}
-	store.cache.Range(func(key, value any) bool {
-		kbz := []byte(key.(string))
+	for key, cv := range store.cache {
+		kbz := []byte(key)
 		if bytes.Compare(kbz, start) < 0 || bytes.Compare(kbz, end) >= 0 {
 			// we don't want to break out of the iteration since cache isn't sorted
-			return true
+			continue
 		}
-		cv := value.(*types.CValue)
 		if cv.Value() == nil {
-			delete(keyStrs, key.(string))
+			delete(keyStrs, key)
 		} else {
-			keyStrs[key.(string)] = struct{}{}
+			keyStrs[key] = struct{}{}
 		}
-		return true
-	})
+	}
 	for k := range keyStrs {
 		res = append(res, k)
 	}