perf: optimize iteration on nested cache context

CHANGELOG.md

@@ -94,6 +94,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#13802](https://github.com/cosmos/cosmos-sdk/pull/13802) Add --output-document flag to the export CLI command to allow writing genesis state to a file.
 * [#13794](https://github.com/cosmos/cosmos-sdk/pull/13794) `types/module.Manager` now supports the
 `cosmossdk.io/core/appmodule.AppModule` API via the new `NewManagerFromMap` constructor.
+* [#13881](https://github.com/cosmos/cosmos-sdk/pull/13881) optimize iteration on nested cache context and other operations in general.
 
 ### State Machine Breaking
 

go.mod

@@ -54,6 +54,7 @@ require (
 	github.com/tendermint/go-amino v0.16.0
 	github.com/tendermint/tendermint v0.37.0-rc1
 	github.com/tendermint/tm-db v0.6.7
+	github.com/tidwall/btree v1.5.2
 	golang.org/x/crypto v0.2.0
 	golang.org/x/exp v0.0.0-20221019170559-20944726eadf
 	google.golang.org/genproto v0.0.0-20221024183307-1bc688fe9f3e

go.sum

@@ -827,6 +827,8 @@ github.com/tendermint/tendermint v0.37.0-rc1 h1:+m+u7s10QD+7vPh5MORrnYjulCdYtGuz
 github.com/tendermint/tendermint v0.37.0-rc1/go.mod h1:z0MZllXL+s0PgIMMpf2P0PrMttQufQio3kUjY2zebeo=
 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.2 h1:5eA83Gfki799V3d3bJo9sWk+yL2LRoTEah3O/SA6/8w=
+github.com/tidwall/btree v1.5.2/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM=
 github.com/ugorji/go v1.2.7 h1:qYhyWUUd6WbiM+C6JZAUkIJt/1WrjzNHY9+KCIjVqTo=

simapp/go.mod

@@ -137,6 +137,7 @@ require (
 	github.com/tendermint/btcd v0.1.1 // indirect
 	github.com/tendermint/crypto v0.0.0-20191022145703-50d29ede1e15 // indirect
 	github.com/tendermint/go-amino v0.16.0 // indirect
+	github.com/tidwall/btree v1.5.2 // indirect
 	github.com/ulikunitz/xz v0.5.8 // indirect
 	github.com/zondax/hid v0.9.1-0.20220302062450-5552068d2266 // indirect
 	go.etcd.io/bbolt v1.3.6 // indirect

simapp/go.sum

@@ -815,6 +815,8 @@ github.com/tendermint/tendermint v0.37.0-rc1 h1:+m+u7s10QD+7vPh5MORrnYjulCdYtGuz
 github.com/tendermint/tendermint v0.37.0-rc1/go.mod h1:z0MZllXL+s0PgIMMpf2P0PrMttQufQio3kUjY2zebeo=
 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.2 h1:5eA83Gfki799V3d3bJo9sWk+yL2LRoTEah3O/SA6/8w=
+github.com/tidwall/btree v1.5.2/go.mod h1:twD9XRA5jj9VUQGELzDO4HPQTNJsoWWfYEL+EUQ2cKY=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM=
 github.com/ugorji/go v1.2.7 h1:qYhyWUUd6WbiM+C6JZAUkIJt/1WrjzNHY9+KCIjVqTo=

store/cachekv/benchmark_test.go

@@ -0,0 +1,162 @@
+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)
+	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.
+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-- {
+		// keep all the cosmos events
+		cs.initialCtx.EventManager().EmitEvents(cs.cachedContexts[i].ctx.EventManager().Events())
+		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))
+	}
+
+	targetCtx := cs.cachedContexts[target].ctx
+	// commit in order from top to bottom
+	for i := len(cs.cachedContexts) - 1; i > target; i-- {
+		// keep all the cosmos events
+		targetCtx.EventManager().EmitEvents(cs.cachedContexts[i].ctx.EventManager().Events())
+		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,85 @@
+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,
+		// we do need to enable locks here, although cachekv is protected with lock,
+		// because multiple iterators could access the btree concurrently.
+		NoLocks: false,
+	})}
+}
+
+func (bt *BTree) Set(key, value []byte) {
+	bt.tree.Set(newPair(key, value))
+}
+
+func (bt *BTree) Get(key []byte) []byte {
+	i, found := bt.tree.Get(newKey(key))
+	if !found {
+		return nil
+	}
+	return i.value
+}
+
+func (bt *BTree) Delete(key []byte) {
+	bt.tree.Delete(newKey(key))
+}
+
+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
+}
+
+// newPair creates a new pair item.
+func newPair(key, value []byte) item {
+	return item{key: key, value: value}
+}
+
+func newKey(key []byte) item {
+	return item{key: key}
+}