feat(mempool): implement RemoveWithReason (#25663)

Author: swift1337

baseapp/abci_utils.go

@@ -285,12 +285,19 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 			return &abci.ResponsePrepareProposal{Txs: h.txSelector.SelectedTxs(ctx)}, nil
 		}
 
-		selectedTxsSignersSeqs := make(map[string]uint64)
+		type invalidTx struct {
+			tx  sdk.Tx
+			err error
+		}
+
 		var (
-			resError        error
-			selectedTxsNums int
-			invalidTxs      []sdk.Tx // invalid txs to be removed out of the loop to avoid dead lock
+			// invalid txs to be removed out of the loop to avoid dead lock
+			invalidTxs             []invalidTx
+			resError               error
+			selectedTxsNums        int
+			selectedTxsSignersSeqs = make(map[string]uint64)
 		)
+
 		mempool.SelectBy(ctx, h.mempool, req.Txs, func(memTx sdk.Tx) bool {
 			unorderedTx, ok := memTx.(sdk.TxWithUnordered)
 			isUnordered := ok && unorderedTx.GetUnordered()
@@ -335,7 +342,7 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 			// check again.
 			txBz, err := h.txVerifier.PrepareProposalVerifyTx(memTx)
 			if err != nil {
-				invalidTxs = append(invalidTxs, memTx)
+				invalidTxs = append(invalidTxs, invalidTx{tx: memTx, err: err})
 			} else {
 				stop := h.txSelector.SelectTxForProposal(ctx, uint64(req.MaxTxBytes), maxBlockGas, memTx, txBz)
 				if stop {
@@ -370,8 +377,13 @@ func (h *DefaultProposalHandler) PrepareProposalHandler() sdk.PrepareProposalHan
 			return nil, resError
 		}
 
-		for _, tx := range invalidTxs {
-			err := h.mempool.Remove(tx)
+		for _, invalidTx := range invalidTxs {
+			reason := mempool.RemoveReason{
+				Caller: mempool.CallerPrepareProposalRemoveInvalid,
+				Error:  invalidTx.err,
+			}
+
+			err := mempool.RemoveWithReason(ctx, h.mempool, invalidTx.tx, reason)
 			if err != nil && !errors.Is(err, mempool.ErrTxNotFound) {
 				return nil, err
 			}

baseapp/baseapp.go

@@ -881,10 +881,16 @@ func (app *BaseApp) RunTx(mode sdk.ExecMode, txBytes []byte, tx sdk.Tx, txIndex
 		if err != nil {
 			if mode == execModeReCheck {
 				// if the ante handler fails on recheck, we want to remove the tx from the mempool
-				if mempoolErr := app.mempool.Remove(tx); mempoolErr != nil {
-					return gInfo, nil, anteEvents, errors.Join(err, mempoolErr)
+				errMempool := mempool.RemoveWithReason(ctx, app.mempool, tx, mempool.RemoveReason{
+					Caller: mempool.CallerRunTxRecheck,
+					Error:  err,
+				})
+
+				if errMempool != nil {
+					return gInfo, nil, anteEvents, errors.Join(err, errMempool)
 				}
 			}
+
 			return gInfo, nil, nil, err
 		}
 
@@ -899,10 +905,10 @@ func (app *BaseApp) RunTx(mode sdk.ExecMode, txBytes []byte, tx sdk.Tx, txIndex
 			return gInfo, nil, anteEvents, err
 		}
 	case execModeFinalize:
-		err = app.mempool.Remove(tx)
+		reason := mempool.RemoveReason{Caller: mempool.CallerRunTxFinalize}
+		err = mempool.RemoveWithReason(ctx, app.mempool, tx, reason)
 		if err != nil && !errors.Is(err, mempool.ErrTxNotFound) {
-			return gInfo, nil, anteEvents,
-				fmt.Errorf("failed to remove tx from mempool: %w", err)
+			return gInfo, nil, anteEvents, fmt.Errorf("failed to remove tx from mempool: %w", err)
 		}
 	}
 

types/mempool/mempool.go

@@ -26,12 +26,34 @@ type Mempool interface {
 
 // ExtMempool is an extension of Mempool interface introduced in v0.50
 // for not be breaking in a patch release.
-// In v0.52+, this interface will be merged into Mempool interface.
+// In v0.54+, this function is removed and RemoveWithReason is merged into Mempool interface.
 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)
+
+	// RemoveWithReason removes a transaction from the mempool with specific reason
+	// allowing the mempool to handle the removal differently.
+	RemoveWithReason(context.Context, sdk.Tx, RemoveReason) error
+}
+
+// RemovalCaller is the origin of the removal
+type RemovalCaller string
+
+// Various callers
+const (
+	CallerRunTxRecheck                 RemovalCaller = "run_tx.recheck"
+	CallerRunTxFinalize                RemovalCaller = "run_tx.finalize"
+	CallerPrepareProposalRemoveInvalid RemovalCaller = "prepare_proposal.remove_invalid"
+)
+
+// RemoveReason is the reason for removing a transaction from the mempool.
+type RemoveReason struct {
+	Caller RemovalCaller
+
+	// Error is an optional error that caused the removal.
+	Error error
 }
 
 // Iterator defines an app-side mempool iterator interface that is as minimal as
@@ -65,3 +87,13 @@ func SelectBy(ctx context.Context, mempool Mempool, txs [][]byte, callback func(
 		iter = iter.Next()
 	}
 }
+
+// RemoveWithReason is compatible with old interface to avoid breaking api.
+// In v0.54+, this function is removed and RemoveWithReason is merged into Mempool interface.
+func RemoveWithReason(ctx context.Context, mempool Mempool, tx sdk.Tx, reason RemoveReason) error {
+	if ext, ok := mempool.(ExtMempool); ok {
+		return ext.RemoveWithReason(ctx, tx, reason)
+	}
+
+	return mempool.Remove(tx)
+}

types/mempool/noop.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) 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) RemoveWithReason(context.Context, sdk.Tx, RemoveReason) error { return nil }

types/mempool/priority_nonce.go

@@ -485,6 +485,11 @@ func (mp *PriorityNonceMempool[C]) Remove(tx sdk.Tx) error {
 	return nil
 }
 
+// RemoveWithReason is a proxy to Remove for this mempool.
+func (mp *PriorityNonceMempool[C]) RemoveWithReason(_ context.Context, tx sdk.Tx, _ RemoveReason) error {
+	return mp.Remove(tx)
+}
+
 func IsEmpty[C comparable](mempool Mempool) error {
 	mp := mempool.(*PriorityNonceMempool[C])
 	if mp.priorityIndex.Len() != 0 {

types/mempool/sender_nonce.go

@@ -254,6 +254,11 @@ func (snm *SenderNonceMempool) Remove(tx sdk.Tx) error {
 	return nil
 }
 
+// RemoveWithReason is a proxy to Remove for this mempool.
+func (snm *SenderNonceMempool) RemoveWithReason(_ context.Context, tx sdk.Tx, _ RemoveReason) error {
+	return snm.Remove(tx)
+}
+
 type senderNonceMempoolIterator struct {
 	rnd           *rand.Rand
 	currentTx     *skiplist.Element