chore: backport pr 13881

CHANGELOG.md

@@ -39,6 +39,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Improvements
 
+* [#13881](https://github.com/cosmos/cosmos-sdk/pull/13881) Optimize iteration on nested cached KV stores and other operations in general.
 * (x/gov) [#14347](https://github.com/cosmos/cosmos-sdk/pull/14347) Support `v1.Proposal` message in `v1beta1.Proposal.Content`.
 
 ### Bug Fixes
@@ -59,13 +60,12 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * (deps) Bump Tendermint version to [v0.34.24](https://github.com/tendermint/tendermint/releases/tag/v0.34.24).
 * [#13651](https://github.com/cosmos/cosmos-sdk/pull/13651) Update `server/config/config.GetConfig` function.
-* [#13781](https://github.com/cosmos/cosmos-sdk/pull/13781) Remove `client/keys.KeysCdc`.
 * [#14175](https://github.com/cosmos/cosmos-sdk/pull/14175) Add `server.DefaultBaseappOptions(appopts)` function to reduce boiler plate in root.go. 
 
 ### State Machine Breaking
 
 * (x/gov) [#14214](https://github.com/cosmos/cosmos-sdk/pull/14214) Fix gov v0.46 migration to v1 votes.
-  * Also provide a helper function `govv046.Migrate_V0466_To_V0467` for migrating a chain already on v0.46 with versions <=v0.46.6 to the latest v0.46.7 correct state.
+    * Also provide a helper function `govv046.Migrate_V0466_To_V0467` for migrating a chain already on v0.46 with versions <=v0.46.6 to the latest v0.46.7 correct state.
 * (x/group) [#14071](https://github.com/cosmos/cosmos-sdk/pull/14071) Don't re-tally proposal after voting period end if they have been marked as ACCEPTED or REJECTED.
 
 ### API Breaking Changes
@@ -239,6 +239,7 @@ replace github.com/confio/ics23/go => github.com/cosmos/cosmos-sdk/ics23/go v0.8
 
 * (x/group) [#12888](https://github.com/cosmos/cosmos-sdk/pull/12888) Fix event propagation to the current context of `x/group` message execution `[]sdk.Result`.
 * (x/upgrade) [#12906](https://github.com/cosmos/cosmos-sdk/pull/12906) Fix upgrade failure by moving downgrade verification logic after store migration.
+* (store) [#12945](https://github.com/cosmos/cosmos-sdk/pull/12945) Fix nil end semantics in store/cachekv/iterator when iterating a dirty cache.
 
 ## [v0.46.0](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.46.0) - 2022-07-26
 

go.mod

@@ -53,6 +53,7 @@ require (
 	github.com/tendermint/go-amino v0.16.0
 	github.com/tendermint/tendermint v0.34.24
 	github.com/tendermint/tm-db v0.6.7
+	github.com/tidwall/btree v1.5.0
 	golang.org/x/crypto v0.2.0
 	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e
 	google.golang.org/genproto v0.0.0-20221014213838-99cd37c6964a

go.sum

@@ -952,6 +952,8 @@ github.com/tendermint/tendermint v0.34.24 h1:879MKKJWYYPJEMMKME+DWUTY4V9f/FBpnZD
 github.com/tendermint/tendermint v0.34.24/go.mod h1:rXVrl4OYzmIa1I91av3iLv2HS0fGSiucyW9J4aMTpKI=
 github.com/tendermint/tm-db v0.6.7 h1:fE00Cbl0jayAoqlExN6oyQJ7fR/ZtoVOmvPJ//+shu8=
 github.com/tendermint/tm-db v0.6.7/go.mod h1:byQDzFkZV1syXr/ReXS808NxA2xvyuuVgXOJ/088L6I=
+github.com/tidwall/btree v1.5.0 h1:iV0yVY/frd7r6qGBXfEYs7DH0gTDgrKTrDjS7xt/IyQ=
+github.com/tidwall/btree v1.5.0/go.mod h1:LGm8L/DZjPLmeWGjv5kFrY8dL4uVhMmzmmLYmsObdKE=
 github.com/tidwall/gjson v1.12.1/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/gjson v1.14.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=

store/cachekv/benchmark_test.go

@@ -0,0 +1,161 @@
+package cachekv_test
+
+import (
+	fmt "fmt"
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/store"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	dbm "github.com/tendermint/tm-db"
+)
+
+func DoBenchmarkDeepContextStack(b *testing.B, depth int) {
+	begin := []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
+	end := []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
+	key := storetypes.NewKVStoreKey("test")
+
+	db := dbm.NewMemDB()
+	cms := store.NewCommitMultiStore(db)
+	cms.MountStoreWithDB(key, storetypes.StoreTypeIAVL, db)
+	cms.LoadLatestVersion()
+	ctx := sdk.NewContext(cms, tmproto.Header{}, false, log.NewNopLogger())
+
+	var stack ContextStack
+	stack.Reset(ctx)
+
+	for i := 0; i < depth; i++ {
+		stack.Snapshot()
+
+		store := stack.CurrentContext().KVStore(key)
+		store.Set(begin, []byte("value"))
+	}
+
+	store := stack.CurrentContext().KVStore(key)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		it := store.Iterator(begin, end)
+		it.Valid()
+		it.Key()
+		it.Value()
+		it.Next()
+		it.Close()
+	}
+}
+
+func BenchmarkDeepContextStack1(b *testing.B) {
+	DoBenchmarkDeepContextStack(b, 1)
+}
+
+func BenchmarkDeepContextStack3(b *testing.B) {
+	DoBenchmarkDeepContextStack(b, 3)
+}
+func BenchmarkDeepContextStack10(b *testing.B) {
+	DoBenchmarkDeepContextStack(b, 10)
+}
+
+func BenchmarkDeepContextStack13(b *testing.B) {
+	DoBenchmarkDeepContextStack(b, 13)
+}
+
+// cachedContext is a pair of cache context and its corresponding commit method.
+// They are obtained from the return value of `context.CacheContext()`.
+type cachedContext struct {
+	ctx    sdk.Context
+	commit func()
+}
+
+// ContextStack manages the initial context and a stack of cached contexts,
+// to support the `StateDB.Snapshot` and `StateDB.RevertToSnapshot` methods.
+//
+// Copied from an old version of ethermint
+type ContextStack struct {
+	// Context of the initial state before transaction execution.
+	// It's the context used by `StateDB.CommitedState`.
+	initialCtx     sdk.Context
+	cachedContexts []cachedContext
+}
+
+// CurrentContext returns the top context of cached stack,
+// if the stack is empty, returns the initial context.
+func (cs *ContextStack) CurrentContext() sdk.Context {
+	l := len(cs.cachedContexts)
+	if l == 0 {
+		return cs.initialCtx
+	}
+	return cs.cachedContexts[l-1].ctx
+}
+
+// Reset sets the initial context and clear the cache context stack.
+func (cs *ContextStack) Reset(ctx sdk.Context) {
+	cs.initialCtx = ctx
+	if len(cs.cachedContexts) > 0 {
+		cs.cachedContexts = []cachedContext{}
+	}
+}
+
+// IsEmpty returns true if the cache context stack is empty.
+func (cs *ContextStack) IsEmpty() bool {
+	return len(cs.cachedContexts) == 0
+}
+
+// Commit commits all the cached contexts from top to bottom in order and clears the stack by setting an empty slice of cache contexts.
+func (cs *ContextStack) Commit() {
+	// commit in order from top to bottom
+	for i := len(cs.cachedContexts) - 1; i >= 0; i-- {
+		if cs.cachedContexts[i].commit == nil {
+			panic(fmt.Sprintf("commit function at index %d should not be nil", i))
+		} else {
+			cs.cachedContexts[i].commit()
+		}
+	}
+	cs.cachedContexts = []cachedContext{}
+}
+
+// CommitToRevision commit the cache after the target revision,
+// to improve efficiency of db operations.
+func (cs *ContextStack) CommitToRevision(target int) error {
+	if target < 0 || target >= len(cs.cachedContexts) {
+		return fmt.Errorf("snapshot index %d out of bound [%d..%d)", target, 0, len(cs.cachedContexts))
+	}
+
+	// commit in order from top to bottom
+	for i := len(cs.cachedContexts) - 1; i > target; i-- {
+		if cs.cachedContexts[i].commit == nil {
+			return fmt.Errorf("commit function at index %d should not be nil", i)
+		}
+		cs.cachedContexts[i].commit()
+	}
+	cs.cachedContexts = cs.cachedContexts[0 : target+1]
+
+	return nil
+}
+
+// Snapshot pushes a new cached context to the stack,
+// and returns the index of it.
+func (cs *ContextStack) Snapshot() int {
+	i := len(cs.cachedContexts)
+	ctx, commit := cs.CurrentContext().CacheContext()
+	cs.cachedContexts = append(cs.cachedContexts, cachedContext{ctx: ctx, commit: commit})
+	return i
+}
+
+// RevertToSnapshot pops all the cached contexts after the target index (inclusive).
+// the target should be snapshot index returned by `Snapshot`.
+// This function panics if the index is out of bounds.
+func (cs *ContextStack) RevertToSnapshot(target int) {
+	if target < 0 || target >= len(cs.cachedContexts) {
+		panic(fmt.Errorf("snapshot index %d out of bound [%d..%d)", target, 0, len(cs.cachedContexts)))
+	}
+	cs.cachedContexts = cs.cachedContexts[:target]
+}
+
+// RevertAll discards all the cache contexts.
+func (cs *ContextStack) RevertAll() {
+	if len(cs.cachedContexts) > 0 {
+		cs.RevertToSnapshot(0)
+	}
+}

store/cachekv/internal/btree.go

@@ -0,0 +1,80 @@
+package internal
+
+import (
+	"bytes"
+	"errors"
+
+	"github.com/tidwall/btree"
+)
+
+const (
+	// The approximate number of items and children per B-tree node. Tuned with benchmarks.
+	// copied from memdb.
+	bTreeDegree = 32
+)
+
+var errKeyEmpty = errors.New("key cannot be empty")
+
+// BTree implements the sorted cache for cachekv store,
+// we don't use MemDB here because cachekv is used extensively in sdk core path,
+// we need it to be as fast as possible, while `MemDB` is mainly used as a mocking db in unit tests.
+//
+// We choose tidwall/btree over google/btree here because it provides API to implement step iterator directly.
+type BTree struct {
+	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 (bt *BTree) Set(key, value []byte) {
+	bt.tree.Set(newItem(key, value))
+}
+
+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) {
+	bt.tree.Delete(newItem(key, nil))
+}
+
+func (bt *BTree) Iterator(start, end []byte) (*memIterator, 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
+}
+
+func (bt *BTree) ReverseIterator(start, end []byte) (*memIterator, 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
+}
+
+// item is a btree item with byte slices as keys and values
+type item struct {
+	key   []byte
+	value []byte
+}
+
+// byKeys compares the items by key
+func byKeys(a, b item) bool {
+	return bytes.Compare(a.key, b.key) == -1
+}
+
+// newItem creates a new pair item.
+func newItem(key, value []byte) item {
+	return item{key: key, value: value}
+}