perf: improve performance of nested cache store

CHANGELOG.md

@@ -225,6 +225,9 @@ extension interfaces. `module.Manager.Modules` is now of type `map[string]interf
 * (snapshots) [14048](https://github.com/cosmos/cosmos-sdk/pull/14048) Move the Snapshot package to the store package. This is done in an effort group all storage related logic under one package.
 * (baseapp) [#14050](https://github.com/cosmos/cosmos-sdk/pull/14050) refactor `ABCIListener` interface to accept go contexts
 * (store/streaming)[#14603](https://github.com/cosmos/cosmos-sdk/pull/14603) `StoreDecoderRegistry` moved from store to `types/simulations` this breaks the `AppModuleSimulation` interface. 
+* (store) [#14444](https://github.com/cosmos/cosmos-sdk/pull/14444) Refactor cache store to improve performance in iteration and nesting cases.
+  - Add new method `CacheMultiStore()` to `Context`, `NewContext` accept `CacheMultiStore` instead of `MultiStore`, change `WithMultiStore` to `WithCacheMultiStore`.
+  - Add new methods `Clone` and `Restore` on `CacheMultiStore` and `CacheKVStore` interfaces.
 
 ### CLI Breaking Changes
 

baseapp/abci_test.go

@@ -866,6 +866,7 @@ func TestABCI_InvalidTransaction(t *testing.T) {
 				require.EqualValues(t, sdkerrors.ErrInvalidSequence.Codespace(), space, err)
 				require.EqualValues(t, sdkerrors.ErrInvalidSequence.ABCICode(), code, err)
 			} else {
+				require.NoError(t, err)
 				require.NotNil(t, result)
 			}
 		}

baseapp/baseapp.go

@@ -9,7 +9,6 @@ import (
 	dbm "github.com/cosmos/cosmos-db"
 	"github.com/cosmos/gogoproto/proto"
 	abci "github.com/tendermint/tendermint/abci/types"
-	"github.com/tendermint/tendermint/crypto/tmhash"
 	"github.com/tendermint/tendermint/libs/log"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	"golang.org/x/exp/maps"
@@ -574,25 +573,6 @@ func (app *BaseApp) getContextForTx(mode runTxMode, txBytes []byte) sdk.Context
 	return ctx
 }
 
-// cacheTxContext returns a new context based off of the provided context with
-// a branched multi-store.
-func (app *BaseApp) cacheTxContext(ctx sdk.Context, txBytes []byte) (sdk.Context, storetypes.CacheMultiStore) {
-	ms := ctx.MultiStore()
-	// TODO: https://github.com/cosmos/cosmos-sdk/issues/2824
-	msCache := ms.CacheMultiStore()
-	if msCache.TracingEnabled() {
-		msCache = msCache.SetTracingContext(
-			storetypes.TraceContext(
-				map[string]interface{}{
-					"txHash": fmt.Sprintf("%X", tmhash.Sum(txBytes)),
-				},
-			),
-		).(storetypes.CacheMultiStore)
-	}
-
-	return ctx.WithMultiStore(msCache), msCache
-}
-
 // runTx processes a transaction within a given execution mode, encoded transaction
 // bytes, and the decoded transaction itself. All state transitions occur through
 // a cached Context depending on the mode provided. State only gets persisted
@@ -607,7 +587,7 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 	var gasWanted uint64
 
 	ctx := app.getContextForTx(mode, txBytes)
-	ms := ctx.MultiStore()
+	cms := ctx.CacheMultiStore()
 
 	// only run the tx if there is block gas remaining
 	if mode == runTxModeDeliver && ctx.BlockGasMeter().IsOutOfGas() {
@@ -655,20 +635,15 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 	}
 
 	if app.anteHandler != nil {
-		var (
-			anteCtx sdk.Context
-			msCache storetypes.CacheMultiStore
-		)
-
 		// Branch context before AnteHandler call in case it aborts.
 		// This is required for both CheckTx and DeliverTx.
 		// Ref: https://github.com/cosmos/cosmos-sdk/issues/2772
 		//
 		// NOTE: Alternatively, we could require that AnteHandler ensures that
 		// writes do not happen if aborted/failed.  This may have some
 		// performance benefits, but it'll be more difficult to get right.
-		anteCtx, msCache = app.cacheTxContext(ctx, txBytes)
-		anteCtx = anteCtx.WithEventManager(sdk.NewEventManager())
+		anteCtx := ctx.WithCacheMultiStore(ctx.CacheMultiStore().Clone()).
+			WithEventManager(sdk.NewEventManager())
 		newCtx, err := app.anteHandler(anteCtx, tx, mode == runTxModeSimulate)
 
 		if !newCtx.IsZero() {
@@ -678,7 +653,7 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 			// Also, in the case of the tx aborting, we need to track gas consumed via
 			// the instantiated gas meter in the AnteHandler, so we update the context
 			// prior to returning.
-			ctx = newCtx.WithMultiStore(ms)
+			ctx = newCtx.WithCacheMultiStore(cms)
 		}
 
 		events := ctx.EventManager().Events()
@@ -690,8 +665,10 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 			return gInfo, nil, nil, 0, err
 		}
 
+		// commit the cloned multistore.
+		cms.Restore(anteCtx.CacheMultiStore())
+
 		priority = ctx.Priority()
-		msCache.Write()
 		anteEvents = events.ToABCIEvents()
 	}
 
@@ -708,17 +685,15 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 		}
 	}
 
-	// Create a new Context based off of the existing Context with a MultiStore branch
-	// in case message processing fails. At this point, the MultiStore
-	// is a branch of a branch.
-	runMsgCtx, msCache := app.cacheTxContext(ctx, txBytes)
-
 	// Attempt to execute all messages and only update state if all messages pass
 	// and we're in DeliverTx. Note, runMsgs will never return a reference to a
 	// Result if any single message fails or does not have a registered Handler.
-	result, err = app.runMsgs(runMsgCtx, msgs, mode)
-
-	if err == nil {
+	err = ctx.RunAtomic(func(runMsgCtx sdk.Context) error {
+		var err error
+		result, err = app.runMsgs(runMsgCtx, msgs, mode)
+		if err != nil {
+			return err
+		}
 
 		// Run optional postHandlers.
 		//
@@ -731,7 +706,8 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 
 			newCtx, err := app.postHandler(postCtx, tx, mode == runTxModeSimulate, err == nil)
 			if err != nil {
-				return gInfo, nil, anteEvents, priority, err
+				result = nil
+				return err
 			}
 
 			result.Events = append(result.Events, newCtx.EventManager().ABCIEvents()...)
@@ -740,15 +716,15 @@ func (app *BaseApp) runTx(mode runTxMode, txBytes []byte) (gInfo sdk.GasInfo, re
 		if mode == runTxModeDeliver {
 			// When block gas exceeds, it'll panic and won't commit the cached store.
 			consumeBlockGas()
-
-			msCache.Write()
 		}
 
 		if len(anteEvents) > 0 && (mode == runTxModeDeliver || mode == runTxModeSimulate) {
 			// append the events in the order of occurrence
 			result.Events = append(anteEvents, result.Events...)
 		}
-	}
+
+		return nil
+	})
 
 	return gInfo, result, anteEvents, priority, err
 }

baseapp/test_helpers.go

@@ -46,8 +46,9 @@ func (app *BaseApp) NewContext(isCheckTx bool, header tmproto.Header) sdk.Contex
 	return sdk.NewContext(app.deliverState.ms, header, false, app.logger)
 }
 
+// NewUncachedContext is used in unit test to bypass the `deliverState.ms`.
 func (app *BaseApp) NewUncachedContext(isCheckTx bool, header tmproto.Header) sdk.Context {
-	return sdk.NewContext(app.cms, header, isCheckTx, app.logger)
+	return sdk.NewContext(app.cms.CacheMultiStore(), header, isCheckTx, app.logger)
 }
 
 func (app *BaseApp) GetContextForDeliverTx(txBytes []byte) sdk.Context {

store/cachekv/benchmark_test.go

@@ -119,7 +119,7 @@ func (cs *CacheStack) CommitToRevision(target int) error {
 // Snapshot pushes a new cached context to the stack,
 // and returns the index of it.
 func (cs *CacheStack) Snapshot() int {
-	cs.cacheStores = append(cs.cacheStores, cachekv.NewStore(cs.CurrentStore()))
+	cs.cacheStores = append(cs.cacheStores, cs.CurrentStore().Clone())
 	return len(cs.cacheStores) - 1
 }
 

store/cachekv/internal/btree.go

@@ -35,16 +35,16 @@ func NewBTree() BTree {
 	}
 }
 
-func (bt BTree) Set(key, value []byte) {
-	bt.tree.Set(newItem(key, value))
+func (bt BTree) Set(key, value []byte, dirty bool) {
+	bt.tree.Set(item{key: key, value: value, dirty: dirty})
 }
 
-func (bt BTree) Get(key []byte) []byte {
+func (bt BTree) Get(key []byte) ([]byte, bool) {
 	i, found := bt.tree.Get(newItem(key, nil))
 	if !found {
-		return nil
+		return nil, false
 	}
-	return i.value
+	return i.value, true
 }
 
 func (bt BTree) Delete(key []byte) {
@@ -65,6 +65,16 @@ func (bt BTree) ReverseIterator(start, end []byte) (types.Iterator, error) {
 	return newMemIterator(start, end, bt, false), nil
 }
 
+// ScanDirtyItems iterate over the dirty entries.
+func (bt BTree) ScanDirtyItems(fn func(key, value []byte)) {
+	bt.tree.Scan(func(item item) bool {
+		if item.dirty {
+			fn(item.key, item.value)
+		}
+		return true
+	})
+}
+
 // 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 {
@@ -77,6 +87,7 @@ func (bt BTree) Copy() BTree {
 type item struct {
 	key   []byte
 	value []byte
+	dirty bool
 }
 
 // byKeys compares the items by key

store/cachekv/internal/btree_test.go

@@ -11,39 +11,44 @@ func TestGetSetDelete(t *testing.T) {
 	db := NewBTree()
 
 	// A nonexistent key should return nil.
-	value := db.Get([]byte("a"))
+	value, found := db.Get([]byte("a"))
 	require.Nil(t, value)
+	require.False(t, found)
 
 	// Set and get a value.
-	db.Set([]byte("a"), []byte{0x01})
-	db.Set([]byte("b"), []byte{0x02})
-	value = db.Get([]byte("a"))
+	db.Set([]byte("a"), []byte{0x01}, true)
+	db.Set([]byte("b"), []byte{0x02}, true)
+	value, found = db.Get([]byte("a"))
 	require.Equal(t, []byte{0x01}, value)
+	require.True(t, found)
 
-	value = db.Get([]byte("b"))
+	value, found = db.Get([]byte("b"))
 	require.Equal(t, []byte{0x02}, value)
+	require.True(t, found)
 
 	// Deleting a non-existent value is fine.
 	db.Delete([]byte("x"))
 
 	// Delete a value.
 	db.Delete([]byte("a"))
 
-	value = db.Get([]byte("a"))
+	value, found = db.Get([]byte("a"))
 	require.Nil(t, value)
+	require.False(t, found)
 
 	db.Delete([]byte("b"))
 
-	value = db.Get([]byte("b"))
+	value, found = db.Get([]byte("b"))
 	require.Nil(t, value)
+	require.False(t, found)
 }
 
 func TestDBIterator(t *testing.T) {
 	db := NewBTree()
 
 	for i := 0; i < 10; i++ {
 		if i != 6 { // but skip 6.
-			db.Set(int642Bytes(int64(i)), []byte{})
+			db.Set(int642Bytes(int64(i)), []byte{}, true)
 		}
 	}