fix: nested iterator on cache store (cosmos#14798)

realiotech · Mar 7, 2023 · 5578d10 · 5578d10
1 parent 7d3d6a7
commit 5578d10
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -39,6 +39,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Bug Fixes
 
+- (store) [#14798](https://github.com/cosmos/cosmos-sdk/pull/14798) Copy btree to avoid the problem of modify while iteration.
 - (cli) [#14799](https://github.com/cosmos/cosmos-sdk/pull/14799) Fix Evidence CLI query flag parsing (backport #13458)
 
 ## [v0.46.8](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.46.8) - 2022-01-23

diff --git a/store/cachekv/internal/btree.go b/store/cachekv/internal/btree.go
@@ -63,6 +63,12 @@ func (bt *BTree) ReverseIterator(start, end []byte) (*memIterator, error) {
 	return NewMemIterator(start, end, bt, make(map[string]struct{}), false), nil
 }
 
+func (bt *BTree) Copy() *BTree {
+	return &BTree{
+		tree: *bt.tree.Copy(),
+	}
+}
+
 // item is a btree item with byte slices as keys and values
 type item struct {
 	key   []byte

diff --git a/store/cachekv/store.go b/store/cachekv/store.go
@@ -188,7 +188,7 @@ func (store *Store) iterator(start, end []byte, ascending bool) types.Iterator {
 	}
 
 	store.dirtyItems(start, end)
-	cache = internal.NewMemIterator(start, end, store.sortedCache, store.deleted, ascending)
+	cache = internal.NewMemIterator(start, end, store.sortedCache.Copy(), store.deleted, ascending)
 
 	return internal.NewCacheMergeIterator(parent, cache, ascending)
 }

diff --git a/store/cachekv/store_test.go b/store/cachekv/store_test.go
@@ -2,6 +2,7 @@ package cachekv_test
 
 import (
 	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -684,3 +685,46 @@ func BenchmarkCacheKVStoreGetKeyFound(b *testing.B) {
 		st.Get([]byte{byte((i & 0xFF0000) >> 16), byte((i & 0xFF00) >> 8), byte(i & 0xFF)})
 	}
 }
+
+func TestIteratorNested(t *testing.T) {
+	mem := dbadapter.Store{DB: dbm.NewMemDB()}
+	store := cachekv.NewStore(mem)
+
+	// setup:
+	// - (owner, contract id) -> contract
+	// - (contract id, record id) -> record
+	owner1 := 1
+	contract1 := 1
+	contract2 := 2
+	record1 := 1
+	record2 := 2
+	store.Set([]byte(fmt.Sprintf("c%04d%04d", owner1, contract1)), []byte("contract1"))
+	store.Set([]byte(fmt.Sprintf("c%04d%04d", owner1, contract2)), []byte("contract2"))
+	store.Set([]byte(fmt.Sprintf("R%04d%04d", contract1, record1)), []byte("contract1-record1"))
+	store.Set([]byte(fmt.Sprintf("R%04d%04d", contract1, record2)), []byte("contract1-record2"))
+	store.Set([]byte(fmt.Sprintf("R%04d%04d", contract2, record1)), []byte("contract2-record1"))
+	store.Set([]byte(fmt.Sprintf("R%04d%04d", contract2, record2)), []byte("contract2-record2"))
+
+	it := types.KVStorePrefixIterator(store, []byte(fmt.Sprintf("c%04d", owner1)))
+	defer it.Close()
+
+	var records []string
+	for ; it.Valid(); it.Next() {
+		contractID, err := strconv.ParseInt(string(it.Key()[5:]), 10, 32)
+		require.NoError(t, err)
+
+		it2 := types.KVStorePrefixIterator(store, []byte(fmt.Sprintf("R%04d", contractID)))
+		for ; it2.Valid(); it2.Next() {
+			records = append(records, string(it2.Value()))
+		}
+
+		it2.Close()
+	}
+
+	require.Equal(t, []string{
+		"contract1-record1",
+		"contract1-record2",
+		"contract2-record1",
+		"contract2-record2",
+	}, records)
+}