feat: improve performance with rpc query (#240)

* feat: improve performance with rpc query

* chore: fix test errors and update go.mod

* add RLock of queryStateMtx in baseapp.Query

baseapp/abci.go

@@ -72,6 +72,7 @@ func (app *BaseApp) InitChain(req abci.RequestInitChain) (res abci.ResponseInitC
 	}
 
 	// initialize states with a correct header
+	app.setQueryState(initHeader)
 	app.setState(runTxModeDeliver, initHeader)
 	app.setState(runTxModeCheck, initHeader)
 
@@ -193,6 +194,18 @@ func (app *BaseApp) BeginBlock(req abci.RequestBeginBlock) (res abci.ResponseBeg
 			WithBlockHeight(req.Header.Height)
 	}
 
+	app.queryStateMtx.Lock()
+	if app.queryState == nil {
+		app.setQueryState(req.Header)
+	} else {
+		// In the first block, app.queryState.ctx will already be initialized
+		// by InitChain. Context is now updated with Header information.
+		app.queryState.ctx = app.queryState.ctx.
+			WithBlockHeader(req.Header).
+			WithBlockHeight(req.Header.Height)
+	}
+	app.queryStateMtx.Unlock()
+
 	gasMeter := app.getBlockGasMeter(app.deliverState.ctx)
 
 	app.deliverState.ctx = app.deliverState.ctx.
@@ -201,9 +214,11 @@ func (app *BaseApp) BeginBlock(req abci.RequestBeginBlock) (res abci.ResponseBeg
 		WithConsensusParams(app.GetConsensusParams(app.deliverState.ctx))
 
 	if app.checkState != nil {
+		app.checkStateMtx.Lock()
 		app.checkState.ctx = app.checkState.ctx.
 			WithBlockGasMeter(gasMeter).
 			WithHeaderHash(req.Hash)
+		app.checkStateMtx.Unlock()
 	}
 
 	if app.beginBlocker != nil {
@@ -381,6 +396,9 @@ func (app *BaseApp) CheckTx(req abci.RequestCheckTx) abci.ResponseCheckTx {
 		panic(fmt.Sprintf("unknown RequestCheckTx type: %s", req.Type))
 	}
 
+	app.checkStateMtx.Lock()
+	defer app.checkStateMtx.Unlock()
+
 	gInfo, result, anteEvents, priority, err := app.runTx(mode, req.Tx)
 	if err != nil {
 		return sdkerrors.ResponseCheckTxWithEvents(err, gInfo.GasWanted, gInfo.GasUsed, anteEvents, app.trace)
@@ -454,9 +472,14 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	// Write the DeliverTx state into branched storage and commit the MultiStore.
 	// The write to the DeliverTx state writes all state transitions to the root
 	// MultiStore (app.cms) so when Commit() is called is persists those values.
+	// checkState needs to be locked here to prevent a race condition
 	app.deliverState.ms.Write()
 	commitID := app.cms.Commit()
 
+	app.queryStateMtx.Lock()
+	app.setQueryState(header)
+	app.queryStateMtx.Unlock()
+
 	res := abci.ResponseCommit{
 		Data:         commitID.Hash,
 		RetainHeight: retainHeight,
@@ -472,10 +495,9 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	app.logger.Info("commit synced", "commit", fmt.Sprintf("%X", commitID))
 
 	// Reset the Check state to the latest committed.
-	//
-	// NOTE: This is safe because Tendermint holds a lock on the mempool for
-	// Commit. Use the header from this latest block.
+	app.checkStateMtx.Lock()
 	app.setState(runTxModeCheck, header)
+	app.checkStateMtx.Unlock()
 
 	// empty/reset the deliver state
 	app.deliverState = nil
@@ -596,7 +618,11 @@ func (app *BaseApp) Query(req abci.RequestQuery) (res abci.ResponseQuery) {
 
 	// when a client did not provide a query height, manually inject the latest
 	if req.Height == 0 {
-		req.Height = app.LastBlockHeight()
+		app.queryStateMtx.RLock()
+		if app.queryState != nil {
+			req.Height = app.queryState.ms.LatestVersion()
+		}
+		app.queryStateMtx.RUnlock()
 	}
 
 	telemetry.IncrCounter(1, "query", "count")
@@ -777,7 +803,12 @@ func (app *BaseApp) ApplySnapshotChunk(req abci.RequestApplySnapshotChunk) abci.
 }
 
 func (app *BaseApp) handleQueryGRPC(handler GRPCQueryHandler, req abci.RequestQuery) abci.ResponseQuery {
-	ctx, err := app.CreateQueryContext(req.Height, req.Prove)
+	if req.Path == "/cosmos.auth.v1beta1.Query/Account" && (req.Height == app.queryState.ms.LatestVersion() || req.Height == 0) {
+		app.checkStateMtx.RLock()
+		defer app.checkStateMtx.RUnlock()
+	}
+
+	ctx, err := app.CreateQueryContext(req.Height, req.Prove, req.Path)
 	if err != nil {
 		return sdkerrors.QueryResult(err, app.trace)
 	}
@@ -793,11 +824,14 @@ func (app *BaseApp) handleQueryGRPC(handler GRPCQueryHandler, req abci.RequestQu
 }
 
 func (app *BaseApp) handleEthQuery(handler EthQueryHandler, req cmtrpctypes.RPCRequest) abci.ResponseEthQuery {
-	// use custom query multistore if provided
-	qms := app.qms
-	if qms == nil {
-		qms = app.cms.(storetypes.MultiStore)
+	// use custom query state if provided
+	app.queryStateMtx.RLock()
+	defer app.queryStateMtx.RUnlock()
+	qs := app.queryState
+	if qs == nil {
+		return sdkerrors.EthQueryResult(fmt.Errorf("queryState is nil"), app.trace)
 	}
+	qms := qs.ms.(sdk.MultiStore)
 
 	height := qms.LatestVersion()
 	if height == 0 {
@@ -815,7 +849,7 @@ func (app *BaseApp) handleEthQuery(handler EthQueryHandler, req cmtrpctypes.RPCR
 	}
 
 	// branch the commit-multistore for safety
-	ctx := sdk.NewContext(cacheMS, app.checkState.ctx.BlockHeader(), true, app.upgradeChecker, app.logger).
+	ctx := sdk.NewContext(cacheMS, qs.ctx.BlockHeader(), true, app.upgradeChecker, app.logger).
 		WithBlockHeight(height)
 
 	res, err := handler(ctx, req)
@@ -860,16 +894,19 @@ func checkNegativeHeight(height int64) error {
 
 // createQueryContext creates a new sdk.Context for a query, taking as args
 // the block height and whether the query needs a proof or not.
-func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, error) {
+func (app *BaseApp) CreateQueryContext(height int64, prove bool, path ...string) (sdk.Context, error) {
 	if err := checkNegativeHeight(height); err != nil {
 		return sdk.Context{}, err
 	}
 
-	// use custom query multistore if provided
-	qms := app.qms
-	if qms == nil {
-		qms = app.cms.(sdk.MultiStore)
+	// use custom query state if provided
+	app.queryStateMtx.RLock()
+	defer app.queryStateMtx.RUnlock()
+	qs := app.queryState
+	if qs == nil {
+		return sdk.Context{}, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, "queryState is nil")
 	}
+	qms := qs.ms.(sdk.MultiStore)
 
 	lastBlockHeight := qms.LatestVersion()
 	if lastBlockHeight == 0 {
@@ -878,9 +915,9 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 
 	if height > lastBlockHeight {
 		return sdk.Context{},
-			sdkerrors.Wrap(
+			sdkerrors.Wrapf(
 				sdkerrors.ErrInvalidHeight,
-				"cannot query with height in the future; please provide a valid height",
+				"cannot query with height in the future(%d, latest height %d); please provide a valid height", height, lastBlockHeight,
 			)
 	}
 
@@ -897,7 +934,15 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 			)
 	}
 
-	cacheMS, err := qms.CacheMultiStoreWithVersion(height)
+	var cacheMS storetypes.CacheMultiStore
+	var err error
+	if len(path) == 1 && path[0] == "/cosmos.auth.v1beta1.Query/Account" && height == lastBlockHeight {
+		// use checkState for account queries on the latest height
+		// we could get the newest account info to send multi txs in one block
+		cacheMS = app.checkState.ms
+	} else {
+		cacheMS, err = qms.CacheMultiStoreWithVersion(height)
+	}
 	if err != nil {
 		return sdk.Context{},
 			sdkerrors.Wrapf(
@@ -907,7 +952,7 @@ func (app *BaseApp) CreateQueryContext(height int64, prove bool) (sdk.Context, e
 	}
 
 	// branch the commit-multistore for safety
-	ctx := sdk.NewContext(cacheMS, app.checkState.ctx.BlockHeader(), true, app.upgradeChecker, app.logger).
+	ctx := sdk.NewContext(cacheMS, qs.ctx.BlockHeader(), true, app.upgradeChecker, app.logger).
 		WithMinGasPrices(app.minGasPrices).
 		WithBlockHeight(height)
 

baseapp/baseapp.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sort"
 	"strings"
+	"sync"
 
 	dbm "github.com/cometbft/cometbft-db"
 	abci "github.com/cometbft/cometbft/abci/types"
@@ -57,7 +58,6 @@ type BaseApp struct { //nolint: maligned
 	name              string               // application name from abci.Info
 	db                dbm.DB               // common DB backend
 	cms               sdk.CommitMultiStore // Main (uncached) state
-	qms               sdk.MultiStore       // Optional alternative multistore for querying only.
 	storeLoader       StoreLoader          // function to handle store loading, may be overridden with SetStoreLoader()
 	grpcQueryRouter   *GRPCQueryRouter     // router for redirecting gRPC query calls
 	ethQueryRouter    *EthQueryRouter      // router for redirecting eth query calls
@@ -85,13 +85,20 @@ type BaseApp struct { //nolint: maligned
 	//
 	// checkState is set on InitChain and reset on Commit
 	// deliverState is set on InitChain and BeginBlock and set to nil on Commit
+	// queryState is set on InitChain and BeginBlock
 	checkState           *state // for CheckTx
 	deliverState         *state // for DeliverTx
 	processProposalState *state // for ProcessProposal
 	prepareProposalState *state // for PrepareProposal
 
 	preDeliverStates []*state // for PreDeliverTx
 
+	// queryState is set on InitChain and BeginBlock
+	queryState *queryState // optional alternative multistore for querying only.
+
+	queryStateMtx sync.RWMutex // mutex for queryState
+	checkStateMtx sync.RWMutex // mutex for checkState
+
 	// an inter-block write-through cache provided to the context during deliverState
 	interBlockCache sdk.MultiStorePersistentCache
 
@@ -180,6 +187,8 @@ func NewBaseApp(
 		txDecoder:        txDecoder,
 		fauxMerkleMode:   false,
 		preDeliverStates: make([]*state, 0),
+		checkStateMtx:    sync.RWMutex{},
+		queryStateMtx:    sync.RWMutex{},
 	}
 
 	for _, option := range options {
@@ -400,6 +409,7 @@ func (app *BaseApp) Init() error {
 	emptyHeader := tmproto.Header{ChainID: app.chainID}
 
 	// needed for the export command which inits from store but never calls initchain
+	app.setQueryState(emptyHeader)
 	app.setState(runTxModeCheck, emptyHeader)
 	app.Seal()
 
@@ -461,8 +471,15 @@ func (app *BaseApp) setState(mode runTxMode, header tmproto.Header) {
 	switch mode {
 	case runTxModeCheck:
 		// Minimum gas prices are also set. It is set on InitChain and reset on Commit.
-		baseState.ctx = baseState.ctx.WithIsCheckTx(true).WithMinGasPrices(app.minGasPrices)
-		app.checkState = baseState
+		var ms sdk.CacheMultiStore
+		if rs, ok := app.cms.(*rootmulti.Store); ok {
+			ms = rs.DeepCopyAndCache()
+			baseState.ms = ms
+			baseState.ctx = baseState.ctx.WithIsCheckTx(true).WithMinGasPrices(app.minGasPrices).WithMultiStore(ms)
+			app.checkState = baseState
+		} else {
+			panic(fmt.Sprintf("set checkState failed: %T is not rootmulti.Store", app.cms))
+		}
 	case runTxModeDeliver:
 		// It is set on InitChain and BeginBlock and set to nil on Commit.
 		app.deliverState = baseState
@@ -483,7 +500,7 @@ func (app *BaseApp) setPreState(number int64, header tmproto.Header) {
 	for i := int64(0); i < number; i++ {
 		var ms sdk.CacheMultiStore
 		if _, ok := app.cms.(*rootmulti.Store); ok {
-			ms = app.cms.(*rootmulti.Store).DeepCopyMultiStore()
+			ms = app.cms.(*rootmulti.Store).DeepCopyAndCache()
 		}
 
 		baseState := &state{
@@ -494,6 +511,19 @@ func (app *BaseApp) setPreState(number int64, header tmproto.Header) {
 	}
 }
 
+func (app *BaseApp) setQueryState(header tmproto.Header) {
+	var ms sdk.CommitMultiStore
+	if rs, ok := app.cms.(*rootmulti.Store); ok {
+		ms = rs.DeepCopy()
+	}
+
+	baseState := &queryState{
+		ms:  ms,
+		ctx: sdk.NewContext(ms, header, false, app.upgradeChecker, app.logger),
+	}
+	app.queryState = baseState
+}
+
 // GetConsensusParams returns the current consensus parameters from the BaseApp's
 // ParamStore. If the BaseApp has no ParamStore defined, nil is returned.
 func (app *BaseApp) GetConsensusParams(ctx sdk.Context) *tmproto.ConsensusParams {

baseapp/options.go

@@ -263,13 +263,6 @@ func (app *BaseApp) SetTxEncoder(txEncoder sdk.TxEncoder) {
 	app.txEncoder = txEncoder
 }
 
-// SetQueryMultiStore set a alternative MultiStore implementation to support grpc query service.
-//
-// Ref: https://github.com/cosmos/cosmos-sdk/issues/13317
-func (app *BaseApp) SetQueryMultiStore(ms sdk.MultiStore) {
-	app.qms = ms
-}
-
 // SetMempool sets the mempool for the BaseApp and is required for the app to start up.
 func (app *BaseApp) SetMempool(mempool mempool.Mempool) {
 	if app.sealed {

baseapp/state.go

@@ -19,3 +19,8 @@ func (st *state) CacheMultiStore() sdk.CacheMultiStore {
 func (st *state) Context() sdk.Context {
 	return st.ctx
 }
+
+type queryState struct {
+	ms  sdk.CommitMultiStore
+	ctx sdk.Context
+}

baseapp/test_helpers.go

@@ -22,6 +22,8 @@ func (app *BaseApp) SimCheck(txEncoder sdk.TxEncoder, tx sdk.Tx) (sdk.GasInfo, *
 
 // Simulate executes a tx in simulate mode to get result and gas info.
 func (app *BaseApp) Simulate(txBytes []byte) (sdk.GasInfo, *sdk.Result, error) {
+	app.checkStateMtx.RLock()
+	defer app.checkStateMtx.RUnlock()
 	gasInfo, result, _, _, err := app.runTx(runTxModeSimulate, txBytes)
 	return gasInfo, result, err
 }

go.mod

@@ -53,7 +53,7 @@ require (
 	github.com/rakyll/statik v0.1.7
 	github.com/rs/zerolog v1.29.1
 	github.com/spf13/cast v1.5.0
-	github.com/spf13/cobra v1.6.1
+	github.com/spf13/cobra v1.7.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.14.0
 	github.com/stretchr/testify v1.8.4
@@ -63,7 +63,7 @@ require (
 	github.com/tidwall/btree v1.6.0
 	github.com/wealdtech/go-eth2-util v1.6.3
 	github.com/willf/bitset v1.1.3
-	golang.org/x/crypto v0.8.0
+	golang.org/x/crypto v0.9.0
 	golang.org/x/exp v0.0.0-20230515195305-f3d0a9c9a5cc
 	golang.org/x/text v0.9.0
 	google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1
@@ -144,7 +144,7 @@ require (
 	github.com/rs/cors v1.8.2 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/sasha-s/go-deadlock v0.3.1 // indirect
-	github.com/sirupsen/logrus v1.9.0 // indirect
+	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/spf13/afero v1.9.2 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/subosito/gotenv v1.4.1 // indirect
@@ -156,9 +156,9 @@ require (
 	github.com/zondax/hid v0.9.1 // indirect
 	github.com/zondax/ledger-go v0.14.1 // indirect
 	go.etcd.io/bbolt v1.3.7 // indirect
-	golang.org/x/net v0.9.0 // indirect
-	golang.org/x/sys v0.7.0 // indirect
-	golang.org/x/term v0.7.0 // indirect
+	golang.org/x/net v0.10.0 // indirect
+	golang.org/x/sys v0.8.0 // indirect
+	golang.org/x/term v0.8.0 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
@@ -171,11 +171,11 @@ replace (
 	// use cosmos fork of keyring
 	github.com/99designs/keyring => github.com/cosmos/keyring v1.2.0
 	github.com/btcsuite/btcd => github.com/btcsuite/btcd v0.23.0
-	// TODO update to official version
-	github.com/cometbft/cometbft => github.com/bnb-chain/greenfield-cometbft v0.0.0-20230705094342-460c09a1e1e0
-	github.com/cometbft/cometbft-db => github.com/bnb-chain/greenfield-cometbft-db v0.0.0-20230706092249-b0e93377dd4d
 
-	github.com/cosmos/iavl => github.com/Pythonberg1997/greenfield-iavl v0.0.0-20230703060649-e52459a6ddca
+	github.com/cometbft/cometbft => github.com/bnb-chain/greenfield-cometbft v0.0.2-alpha.1
+	github.com/cometbft/cometbft-db => github.com/bnb-chain/greenfield-cometbft-db v0.8.1-alpha.1
+	github.com/cosmos/iavl => github.com/bnb-chain/greenfield-iavl v0.20.1-alpha.1
+
 	// Downgraded to avoid bugs in following commits which caused simulations to fail.
 	// dgrijalva/jwt-go is deprecated and doesn't receive security updates.
 	// TODO: remove it: https://github.com/cosmos/cosmos-sdk/issues/13134