refactor: make cachekv store thread-safe again (backport #14378)

CHANGELOG.md

@@ -51,6 +51,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### State Machine Breaking
 
 * (baseapp, x/auth/posthandler) [#13940](https://github.com/cosmos/cosmos-sdk/pull/13940) Update `PostHandler` to receive the `runTx` success boolean.
+* (store) [#14378](https://github.com/cosmos/cosmos-sdk/pull/14378) The `CacheKV` store is thread-safe again, which includes improved iteration and deletion logic. Iteration is on a strictly isolated view now, which is breaking from previous behavior.
 
 ### API Breaking Changes
 

store/cachekv/internal/btree.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 
+	"github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/tidwall/btree"
 )
 
@@ -21,46 +22,55 @@ var errKeyEmpty = errors.New("key cannot be empty")
 //
 // We choose tidwall/btree over google/btree here because it provides API to implement step iterator directly.
 type BTree struct {
-	tree btree.BTreeG[item]
+	tree *btree.BTreeG[item]
 }
 
 // NewBTree creates a wrapper around `btree.BTreeG`.
-func NewBTree() *BTree {
-	return &BTree{tree: *btree.NewBTreeGOptions(byKeys, btree.Options{
-		Degree: bTreeDegree,
-		// Contract: cachekv store must not be called concurrently
-		NoLocks: true,
-	})}
+func NewBTree() BTree {
+	return BTree{
+		tree: btree.NewBTreeGOptions(byKeys, btree.Options{
+			Degree:  bTreeDegree,
+			NoLocks: false,
+		}),
+	}
 }
 
-func (bt *BTree) Set(key, value []byte) {
+func (bt BTree) Set(key, value []byte) {
 	bt.tree.Set(newItem(key, value))
 }
 
-func (bt *BTree) Get(key []byte) []byte {
+func (bt BTree) Get(key []byte) []byte {
 	i, found := bt.tree.Get(newItem(key, nil))
 	if !found {
 		return nil
 	}
 	return i.value
 }
 
-func (bt *BTree) Delete(key []byte) {
+func (bt BTree) Delete(key []byte) {
 	bt.tree.Delete(newItem(key, nil))
 }
 
-func (bt *BTree) Iterator(start, end []byte) (*memIterator, error) {
+func (bt BTree) Iterator(start, end []byte) (types.Iterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {
 		return nil, errKeyEmpty
 	}
-	return NewMemIterator(start, end, bt, make(map[string]struct{}), true), nil
+	return newMemIterator(start, end, bt, true), nil
 }
 
-func (bt *BTree) ReverseIterator(start, end []byte) (*memIterator, error) {
+func (bt BTree) ReverseIterator(start, end []byte) (types.Iterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {
 		return nil, errKeyEmpty
 	}
-	return NewMemIterator(start, end, bt, make(map[string]struct{}), false), nil
+	return newMemIterator(start, end, bt, false), nil
+}
+
+// Copy the tree. This is a copy-on-write operation and is very fast because
+// it only performs a shadowed copy.
+func (bt BTree) Copy() BTree {
+	return BTree{
+		tree: bt.tree.Copy(),
+	}
 }
 
 // item is a btree item with byte slices as keys and values

store/cachekv/internal/btree_test.go

@@ -3,6 +3,7 @@ package internal
 import (
 	"testing"
 
+	"github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/stretchr/testify/require"
 )
@@ -181,7 +182,7 @@ func TestDBIterator(t *testing.T) {
 	verifyIterator(t, ritr, nil, "reverse iterator with empty db")
 }
 
-func verifyIterator(t *testing.T, itr *memIterator, expected []int64, msg string) {
+func verifyIterator(t *testing.T, itr types.Iterator, expected []int64, msg string) {
 	i := 0
 	for itr.Valid() {
 		key := itr.Key()

store/cachekv/internal/memiterator.go

@@ -11,20 +11,18 @@ import (
 var _ types.Iterator = (*memIterator)(nil)
 
 // memIterator iterates over iterKVCache items.
-// if key is nil, means it was deleted.
+// if value is nil, means it was deleted.
 // Implements Iterator.
 type memIterator struct {
 	iter btree.GenericIter[item]
 
 	start     []byte
 	end       []byte
 	ascending bool
-	lastKey   []byte
-	deleted   map[string]struct{}
 	valid     bool
 }
 
-func NewMemIterator(start, end []byte, items *BTree, deleted map[string]struct{}, ascending bool) *memIterator {
+func newMemIterator(start, end []byte, items BTree, ascending bool) *memIterator {
 	iter := items.tree.Iter()
 	var valid bool
 	if ascending {
@@ -52,8 +50,6 @@ func NewMemIterator(start, end []byte, items *BTree, deleted map[string]struct{}
 		start:     start,
 		end:       end,
 		ascending: ascending,
-		lastKey:   nil,
-		deleted:   deleted,
 		valid:     valid,
 	}
 
@@ -113,21 +109,7 @@ func (mi *memIterator) Key() []byte {
 }
 
 func (mi *memIterator) Value() []byte {
-	item := mi.iter.Item()
-	key := item.key
-	// 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.
-	if _, ok := mi.deleted[string(key)]; ok {
-		if mi.lastKey == nil || !bytes.Equal(key, mi.lastKey) {
-			// not re-calling on old last key
-			return nil
-		}
-	}
-	mi.lastKey = key
-	return item.value
+	return mi.iter.Item().value
 }
 
 func (mi *memIterator) assertValid() {

store/cachekv/store.go

@@ -24,14 +24,11 @@ type cValue struct {
 }
 
 // Store wraps an in-memory cache around an underlying types.KVStore.
-// If a cached value is nil but deleted is defined for the corresponding key,
-// it means the parent doesn't have the key. (No need to delete upon Write())
 type Store struct {
 	mtx           sync.Mutex
 	cache         map[string]*cValue
-	deleted       map[string]struct{}
 	unsortedCache map[string]struct{}
-	sortedCache   *internal.BTree // always ascending sorted
+	sortedCache   internal.BTree // always ascending sorted
 	parent        types.KVStore
 }
 
@@ -41,7 +38,6 @@ var _ types.CacheKVStore = (*Store)(nil)
 func NewStore(parent types.KVStore) *Store {
 	return &Store{
 		cache:         make(map[string]*cValue),
-		deleted:       make(map[string]struct{}),
 		unsortedCache: make(map[string]struct{}),
 		sortedCache:   internal.NewBTree(),
 		parent:        parent,
@@ -63,7 +59,7 @@ func (store *Store) Get(key []byte) (value []byte) {
 	cacheValue, ok := store.cache[conv.UnsafeBytesToStr(key)]
 	if !ok {
 		value = store.parent.Get(key)
-		store.setCacheValue(key, value, false, false)
+		store.setCacheValue(key, value, false)
 	} else {
 		value = cacheValue.value
 	}
@@ -79,7 +75,7 @@ func (store *Store) Set(key []byte, value []byte) {
 	types.AssertValidKey(key)
 	types.AssertValidValue(value)
 
-	store.setCacheValue(key, value, false, true)
+	store.setCacheValue(key, value, true)
 }
 
 // Has implements types.KVStore.
@@ -94,15 +90,15 @@ func (store *Store) Delete(key []byte) {
 	defer store.mtx.Unlock()
 
 	types.AssertValidKey(key)
-	store.setCacheValue(key, nil, true, true)
+	store.setCacheValue(key, nil, true)
 }
 
 // Implements Cachetypes.KVStore.
 func (store *Store) Write() {
 	store.mtx.Lock()
 	defer store.mtx.Unlock()
 
-	if len(store.cache) == 0 && len(store.deleted) == 0 && len(store.unsortedCache) == 0 {
+	if len(store.cache) == 0 && len(store.unsortedCache) == 0 {
 		store.sortedCache = internal.NewBTree()
 		return
 	}
@@ -122,19 +118,16 @@ func (store *Store) Write() {
 	// TODO: Consider allowing usage of Batch, which would allow the write to
 	// at least happen atomically.
 	for _, key := range keys {
-		if store.isDeleted(key) {
-			// We use []byte(key) instead of conv.UnsafeStrToBytes because we cannot
-			// be sure if the underlying store might do a save with the byteslice or
-			// not. Once we get confirmation that .Delete is guaranteed not to
-			// save the byteslice, then we can assume only a read-only copy is sufficient.
-			store.parent.Delete([]byte(key))
-			continue
-		}
-
+		// We use []byte(key) instead of conv.UnsafeStrToBytes because we cannot
+		// be sure if the underlying store might do a save with the byteslice or
+		// not. Once we get confirmation that .Delete is guaranteed not to
+		// save the byteslice, then we can assume only a read-only copy is sufficient.
 		cacheValue := store.cache[key]
 		if cacheValue.value != nil {
-			// It already exists in the parent, hence delete it.
+			// It already exists in the parent, hence update it.
 			store.parent.Set([]byte(key), cacheValue.value)
+		} else {
+			store.parent.Delete([]byte(key))
 		}
 	}
 
@@ -144,9 +137,6 @@ func (store *Store) Write() {
 	for key := range store.cache {
 		delete(store.cache, key)
 	}
-	for key := range store.deleted {
-		delete(store.deleted, key)
-	}
 	for key := range store.unsortedCache {
 		delete(store.unsortedCache, key)
 	}
@@ -180,16 +170,24 @@ func (store *Store) iterator(start, end []byte, ascending bool) types.Iterator {
 	store.mtx.Lock()
 	defer store.mtx.Unlock()
 
-	var parent, cache types.Iterator
+	store.dirtyItems(start, end)
+	isoSortedCache := store.sortedCache.Copy()
+
+	var (
+		err           error
+		parent, cache types.Iterator
+	)
 
 	if ascending {
 		parent = store.parent.Iterator(start, end)
+		cache, err = isoSortedCache.Iterator(start, end)
 	} else {
 		parent = store.parent.ReverseIterator(start, end)
+		cache, err = isoSortedCache.ReverseIterator(start, end)
+	}
+	if err != nil {
+		panic(err)
 	}
-
-	store.dirtyItems(start, end)
-	cache = internal.NewMemIterator(start, end, store.sortedCache, store.deleted, ascending)
 
 	return internal.NewCacheMergeIterator(parent, cache, ascending)
 }
@@ -370,13 +368,7 @@ func (store *Store) clearUnsortedCacheSubset(unsorted []*kv.Pair, sortState sort
 	}
 
 	for _, item := range unsorted {
-		if item.Value == nil {
-			// deleted element, tracked by store.deleted
-			// setting arbitrary value
-			store.sortedCache.Set(item.Key, []byte{})
-			continue
-		}
-
+		// sortedCache is able to store `nil` value to represent deleted items.
 		store.sortedCache.Set(item.Key, item.Value)
 	}
 }
@@ -385,23 +377,14 @@ func (store *Store) clearUnsortedCacheSubset(unsorted []*kv.Pair, sortState sort
 // etc
 
 // Only entrypoint to mutate store.cache.
-func (store *Store) setCacheValue(key, value []byte, deleted bool, dirty bool) {
+// A `nil` value means a deletion.
+func (store *Store) setCacheValue(key, value []byte, dirty bool) {
 	keyStr := conv.UnsafeBytesToStr(key)
 	store.cache[keyStr] = &cValue{
 		value: value,
 		dirty: dirty,
 	}
-	if deleted {
-		store.deleted[keyStr] = struct{}{}
-	} else {
-		delete(store.deleted, keyStr)
-	}
 	if dirty {
 		store.unsortedCache[keyStr] = struct{}{}
 	}
 }
-
-func (store *Store) isDeleted(key string) bool {
-	_, ok := store.deleted[key]
-	return ok
-}