feat(mempool): respect gas wanted returned by ante handler

* Problem: mempool don't respect gas wanted returned by ante handler

Solution:
- support custom gas wanted in mempool

* Update CHANGELOG.md

Signed-off-by: yihuang <[email protected]>

* cleanup

* cleanup

* fix priorityIndex

* fix process proposal

* fix lint

* fix lint

---------

Signed-off-by: yihuang <[email protected]>

Problem: mempool.GasTx interface is not reused (#536)

* Problem: redundant mutex for InsertWithGasWanted

cfg of PriorityNonceMempool remains unchanged once assigned, so no lock is required

* make mocks

* cleanup

* keep order of check MaxTx

Author: yihuang

CHANGELOG.md

@@ -62,6 +62,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (baseapp) [#24655](https://github.com/cosmos/cosmos-sdk/pull/24655) Add mutex locks for `state` and make `lastCommitInfo` atomic to prevent race conditions between `Commit` and `CreateQueryContext`.
 * (proto) [#24161](https://github.com/cosmos/cosmos-sdk/pull/24161) Remove unnecessary annotations from `x/staking` authz proto.
 * (x/bank) [#24660](https://github.com/cosmos/cosmos-sdk/pull/24660) Improve performance of the `GetAllBalances` and `GetAccountsBalances` keeper methods.
+* (mempool)[#25338](https://github.com/cosmos/cosmos-sdk/pull/25338) Respect gas wanted returned by ante handler for mempool.
 
 ### Bug Fixes
 

baseapp/abci_utils.go

@@ -201,7 +201,7 @@ type (
 	// to verify a transaction.
 	ProposalTxVerifier interface {
 		PrepareProposalVerifyTx(tx sdk.Tx) ([]byte, error)
-		ProcessProposalVerifyTx(txBz []byte) (sdk.Tx, error)
+		ProcessProposalVerifyTx(txBz []byte) (sdk.Tx, uint64, error)
 		TxDecode(txBz []byte) (sdk.Tx, error)
 		TxEncode(tx sdk.Tx) ([]byte, error)
 	}
@@ -276,7 +276,12 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 					return nil, err
 				}
 
-				stop := h.txSelector.SelectTxForProposal(ctx, uint64(req.MaxTxBytes), maxBlockGas, tx, txBz)
+				var txGasLimit uint64
+				if gasTx, ok := tx.(mempool.GasTx); ok {
+					txGasLimit = gasTx.GetGas()
+				}
+
+				stop := h.txSelector.SelectTxForProposal(ctx, uint64(req.MaxTxBytes), maxBlockGas, tx, txBz, txGasLimit)
 				if stop {
 					break
 				}
@@ -291,14 +296,14 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 			selectedTxsNums int
 			invalidTxs      []sdk.Tx // invalid txs to be removed out of the loop to avoid dead lock
 		)
-		mempool.SelectBy(ctx, h.mempool, req.Txs, func(memTx sdk.Tx) bool {
-			unorderedTx, ok := memTx.(sdk.TxWithUnordered)
+		mempool.SelectBy(ctx, h.mempool, req.Txs, func(memTx mempool.Tx) bool {
+			unorderedTx, ok := memTx.Tx.(sdk.TxWithUnordered)
 			isUnordered := ok && unorderedTx.GetUnordered()
 			txSignersSeqs := make(map[string]uint64)
 
 			// if the tx is unordered, we don't need to check the sequence, we just add it
 			if !isUnordered {
-				signerData, err := h.signerExtAdapter.GetSigners(memTx)
+				signerData, err := h.signerExtAdapter.GetSigners(memTx.Tx)
 				if err != nil {
 					// propagate the error to the caller
 					resError = err
@@ -333,11 +338,11 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 			// which calls mempool.Insert, in theory everything in the pool should be
 			// valid. But some mempool implementations may insert invalid txs, so we
 			// check again.
-			txBz, err := h.txVerifier.PrepareProposalVerifyTx(memTx)
+			txBz, err := h.txVerifier.PrepareProposalVerifyTx(memTx.Tx)
 			if err != nil {
-				invalidTxs = append(invalidTxs, memTx)
+				invalidTxs = append(invalidTxs, memTx.Tx)
 			} else {
-				stop := h.txSelector.SelectTxForProposal(ctx, uint64(req.MaxTxBytes), maxBlockGas, memTx, txBz)
+				stop := h.txSelector.SelectTxForProposal(ctx, uint64(req.MaxTxBytes), maxBlockGas, memTx.Tx, txBz, memTx.GasWanted)
 				if stop {
 					return false
 				}
@@ -409,17 +414,13 @@ func (h *DefaultProposalHandler) ProcessProposalHandler() sdk.ProcessProposalHan
 		}
 
 		for _, txBytes := range req.Txs {
-			tx, err := h.txVerifier.ProcessProposalVerifyTx(txBytes)
+			_, gasWanted, err := h.txVerifier.ProcessProposalVerifyTx(txBytes)
 			if err != nil {
 				return &abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_REJECT}, nil
 			}
 
 			if maxBlockGas > 0 {
-				gasTx, ok := tx.(GasTx)
-				if ok {
-					totalTxGas += gasTx.GetGas()
-				}
-
+				totalTxGas += gasWanted
 				if totalTxGas > uint64(maxBlockGas) {
 					return &abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_REJECT}, nil
 				}
@@ -477,7 +478,7 @@ type TxSelector interface {
 	// a proposal based on inclusion criteria defined by the TxSelector. It must
 	// return <true> if the caller should halt the transaction selection loop
 	// (typically over a mempool) or <false> otherwise.
-	SelectTxForProposal(ctx context.Context, maxTxBytes, maxBlockGas uint64, memTx sdk.Tx, txBz []byte) bool
+	SelectTxForProposal(ctx context.Context, maxTxBytes, maxBlockGas uint64, memTx sdk.Tx, txBz []byte, gasWanted uint64) bool
 }
 
 type defaultTxSelector struct {
@@ -502,23 +503,16 @@ func (ts *defaultTxSelector) Clear() {
 	ts.selectedTxs = nil
 }
 
-func (ts *defaultTxSelector) SelectTxForProposal(_ context.Context, maxTxBytes, maxBlockGas uint64, memTx sdk.Tx, txBz []byte) bool {
+func (ts *defaultTxSelector) SelectTxForProposal(_ context.Context, maxTxBytes, maxBlockGas uint64, memTx sdk.Tx, txBz []byte, gasWanted uint64) bool {
 	txSize := uint64(cmttypes.ComputeProtoSizeForTxs([]cmttypes.Tx{txBz}))
 
-	var txGasLimit uint64
-	if memTx != nil {
-		if gasTx, ok := memTx.(GasTx); ok {
-			txGasLimit = gasTx.GetGas()
-		}
-	}
-
 	// only add the transaction to the proposal if we have enough capacity
 	if (txSize + ts.totalTxBytes) <= maxTxBytes {
 		// If there is a max block gas limit, add the tx only if the limit has
 		// not been met.
 		if maxBlockGas > 0 {
-			if (txGasLimit + ts.totalTxGas) <= maxBlockGas {
-				ts.totalTxGas += txGasLimit
+			if (gasWanted + ts.totalTxGas) <= maxBlockGas {
+				ts.totalTxGas += gasWanted
 				ts.totalTxBytes += txSize
 				ts.selectedTxs = append(ts.selectedTxs, txBz)
 			}

baseapp/baseapp.go

@@ -873,7 +873,7 @@ func (app *BaseApp) runTx(mode sdk.ExecMode, txBytes []byte, tx sdk.Tx) (gInfo s
 
 	switch mode {
 	case execModeCheck:
-		err = app.mempool.Insert(ctx, tx)
+		err = app.mempool.InsertWithGasWanted(ctx, tx, gasWanted)
 		if err != nil {
 			return gInfo, nil, anteEvents, err
 		}
@@ -1067,18 +1067,18 @@ func (app *BaseApp) PrepareProposalVerifyTx(tx sdk.Tx) ([]byte, error) {
 // ProcessProposal state internally will be discarded. <nil, err> will be
 // returned if the transaction cannot be decoded. <Tx, nil> will be returned if
 // the transaction is valid, otherwise <Tx, err> will be returned.
-func (app *BaseApp) ProcessProposalVerifyTx(txBz []byte) (sdk.Tx, error) {
+func (app *BaseApp) ProcessProposalVerifyTx(txBz []byte) (sdk.Tx, uint64, error) {
 	tx, err := app.txDecoder(txBz)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	_, _, _, err = app.runTx(execModeProcessProposal, txBz, tx)
+	gInfo, _, _, err := app.runTx(execModeProcessProposal, txBz, tx)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
-	return tx, nil
+	return tx, gInfo.GasWanted, nil
 }
 
 func (app *BaseApp) TxDecode(txBytes []byte) (sdk.Tx, error) {

baseapp/testutil/mock/mocks.go

@@ -7,11 +7,30 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
+type Tx struct {
+	Tx        sdk.Tx
+	GasWanted uint64
+}
+
+func NewMempoolTx(tx sdk.Tx, gasWanted uint64) Tx {
+	return Tx{
+		Tx:        tx,
+		GasWanted: gasWanted,
+	}
+}
+
+type GasTx interface {
+	GetGas() uint64
+}
+
 type Mempool interface {
 	// Insert attempts to insert a Tx into the app-side mempool returning
 	// an error upon failure.
 	Insert(context.Context, sdk.Tx) error
 
+	// Insert with a custom gas wanted value
+	InsertWithGasWanted(context.Context, sdk.Tx, uint64) error
+
 	// Select returns an Iterator over the app-side mempool. If txs are specified,
 	// then they shall be incorporated into the Iterator. The Iterator is not thread-safe to use.
 	Select(context.Context, [][]byte) Iterator
@@ -31,7 +50,7 @@ type ExtMempool interface {
 	Mempool
 
 	// SelectBy use callback to iterate over the mempool, it's thread-safe to use.
-	SelectBy(context.Context, [][]byte, func(sdk.Tx) bool)
+	SelectBy(context.Context, [][]byte, func(Tx) bool)
 }
 
 // Iterator defines an app-side mempool iterator interface that is as minimal as
@@ -43,7 +62,7 @@ type Iterator interface {
 	Next() Iterator
 
 	// Tx returns the transaction at the current position of the iterator.
-	Tx() sdk.Tx
+	Tx() Tx
 }
 
 var (
@@ -53,7 +72,7 @@ var (
 
 // SelectBy is compatible with old interface to avoid breaking api.
 // In v0.52+, this function is removed and SelectBy is merged into Mempool interface.
-func SelectBy(ctx context.Context, mempool Mempool, txs [][]byte, callback func(sdk.Tx) bool) {
+func SelectBy(ctx context.Context, mempool Mempool, txs [][]byte, callback func(Tx) bool) {
 	if ext, ok := mempool.(ExtMempool); ok {
 		ext.SelectBy(ctx, txs, callback)
 		return

types/mempool/mempool.go

@@ -139,7 +139,7 @@ func fetchTxs(iterator mempool.Iterator, maxBytes int64) []sdk.Tx {
 		if numBytes += txSize; numBytes > maxBytes {
 			break
 		}
-		txs = append(txs, iterator.Tx())
+		txs = append(txs, iterator.Tx().Tx)
 		i := iterator.Next()
 		iterator = i
 	}

types/mempool/mempool_test.go

@@ -16,8 +16,9 @@ var _ ExtMempool = (*NoOpMempool)(nil)
 // is FIFO-ordered by default.
 type NoOpMempool struct{}
 
-func (NoOpMempool) Insert(context.Context, sdk.Tx) error                  { return nil }
-func (NoOpMempool) Select(context.Context, [][]byte) Iterator             { return nil }
-func (NoOpMempool) SelectBy(context.Context, [][]byte, func(sdk.Tx) bool) {}
-func (NoOpMempool) CountTx() int                                          { return 0 }
-func (NoOpMempool) Remove(sdk.Tx) error                                   { return nil }
+func (NoOpMempool) Insert(context.Context, sdk.Tx) error                      { return nil }
+func (NoOpMempool) InsertWithGasWanted(context.Context, sdk.Tx, uint64) error { return nil }
+func (NoOpMempool) Select(context.Context, [][]byte) Iterator                 { return nil }
+func (NoOpMempool) SelectBy(context.Context, [][]byte, func(Tx) bool)         {}
+func (NoOpMempool) CountTx() int                                              { return 0 }
+func (NoOpMempool) Remove(sdk.Tx) error                                       { return nil }

types/mempool/noop.go

@@ -189,10 +189,10 @@ func (mp *PriorityNonceMempool[C]) NextSenderTx(sender string) sdk.Tx {
 	}
 
 	cursor := senderIndex.Front()
-	return cursor.Value.(sdk.Tx)
+	return cursor.Value.(Tx).Tx
 }
 
-// Insert attempts to insert a Tx into the app-side mempool in O(log n) time,
+// InsertWithGasWanted attempts to insert a Tx into the app-side mempool in O(log n) time,
 // returning an error if unsuccessful. Sender and nonce are derived from the
 // transaction's first signature.
 //
@@ -201,7 +201,7 @@ func (mp *PriorityNonceMempool[C]) NextSenderTx(sender string) sdk.Tx {
 //
 // Inserting a duplicate tx with a different priority overwrites the existing tx,
 // changing the total order of the mempool.
-func (mp *PriorityNonceMempool[C]) Insert(ctx context.Context, tx sdk.Tx) error {
+func (mp *PriorityNonceMempool[C]) InsertWithGasWanted(ctx context.Context, tx sdk.Tx, gasWanted uint64) error {
 	mp.mtx.Lock()
 	defer mp.mtx.Unlock()
 	if mp.cfg.MaxTx > 0 && mp.priorityIndex.Len() >= mp.cfg.MaxTx {
@@ -210,6 +210,8 @@ func (mp *PriorityNonceMempool[C]) Insert(ctx context.Context, tx sdk.Tx) error
 		return nil
 	}
 
+	memTx := NewMempoolTx(tx, gasWanted)
+
 	sigs, err := mp.cfg.SignerExtractor.GetSigners(tx)
 	if err != nil {
 		return err
@@ -247,12 +249,12 @@ func (mp *PriorityNonceMempool[C]) Insert(ctx context.Context, tx sdk.Tx) error
 	// changes.
 	sk := txMeta[C]{nonce: nonce, sender: sender}
 	if oldScore, txExists := mp.scores[sk]; txExists {
-		if mp.cfg.TxReplacement != nil && !mp.cfg.TxReplacement(oldScore.priority, priority, senderIndex.Get(key).Value.(sdk.Tx), tx) {
+		if mp.cfg.TxReplacement != nil && !mp.cfg.TxReplacement(oldScore.priority, priority, senderIndex.Get(key).Value.(Tx).Tx, tx) {
 			return fmt.Errorf(
 				"tx doesn't fit the replacement rule, oldPriority: %v, newPriority: %v, oldTx: %v, newTx: %v",
 				oldScore.priority,
 				priority,
-				senderIndex.Get(key).Value.(sdk.Tx),
+				senderIndex.Get(key).Value.(Tx).Tx,
 				tx,
 			)
 		}
@@ -270,14 +272,23 @@ func (mp *PriorityNonceMempool[C]) Insert(ctx context.Context, tx sdk.Tx) error
 
 	// Since senderIndex is scored by nonce, a changed priority will overwrite the
 	// existing key.
-	key.senderElement = senderIndex.Set(key, tx)
+	key.senderElement = senderIndex.Set(key, memTx)
 
 	mp.scores[sk] = txMeta[C]{priority: priority}
 	mp.priorityIndex.Set(key, tx)
 
 	return nil
 }
 
+func (mp *PriorityNonceMempool[C]) Insert(ctx context.Context, tx sdk.Tx) error {
+	var gasLimit uint64
+	if gasTx, ok := tx.(GasTx); ok {
+		gasLimit = gasTx.GetGas()
+	}
+
+	return mp.InsertWithGasWanted(ctx, tx, gasLimit)
+}
+
 func (i *PriorityNonceIterator[C]) iteratePriority() Iterator {
 	// beginning of priority iteration
 	if i.priorityNode == nil {
@@ -341,8 +352,8 @@ func (i *PriorityNonceIterator[C]) Next() Iterator {
 	return i
 }
 
-func (i *PriorityNonceIterator[C]) Tx() sdk.Tx {
-	return i.senderCursors[i.sender].Value.(sdk.Tx)
+func (i *PriorityNonceIterator[C]) Tx() Tx {
+	return i.senderCursors[i.sender].Value.(Tx)
 }
 
 // Select returns a set of transactions from the mempool, ordered by priority
@@ -376,7 +387,7 @@ func (mp *PriorityNonceMempool[C]) doSelect(_ context.Context, _ [][]byte) Itera
 }
 
 // SelectBy will hold the mutex during the iteration, callback returns if continue.
-func (mp *PriorityNonceMempool[C]) SelectBy(ctx context.Context, txs [][]byte, callback func(sdk.Tx) bool) {
+func (mp *PriorityNonceMempool[C]) SelectBy(ctx context.Context, txs [][]byte, callback func(Tx) bool) {
 	mp.mtx.Lock()
 	defer mp.mtx.Unlock()