Made autobahn producer use TxMempool

sei-tendermint/internal/autobahn/producer/state.go

@@ -6,8 +6,8 @@ import (
 	"time"
 
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/autobahn/consensus"
-	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/autobahn/pb"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/autobahn/types"
+	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/mempool"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils/scope"
 	"golang.org/x/time/rate"
@@ -32,19 +32,18 @@ func (c *Config) maxTxsPerBlock() uint64 {
 
 // State is the block producer state.
 type State struct {
-	cfg *Config
-	// channel of transactions to build the next block from.
-	mempool chan *pb.Transaction
+	cfg       *Config
+	txMempool *mempool.TxMempool
 	// consensus state to which published blocks will be reported.
 	consensus *consensus.State
 }
 
 // NewState constructs a new block producer state.
 // Returns an error if the current node is NOT a producer.
-func NewState(cfg *Config, consensus *consensus.State) *State {
+func NewState(cfg *Config, txMempool *mempool.TxMempool, consensus *consensus.State) *State {
 	return &State{
 		cfg:       cfg,
-		mempool:   make(chan *pb.Transaction, cfg.MempoolSize),
+		txMempool: txMempool,
 		consensus: consensus,
 	}
 }
@@ -54,21 +53,32 @@ func NewState(cfg *Config, consensus *consensus.State) *State {
 func (s *State) makePayload(ctx context.Context) *types.Payload {
 	ctx, cancel := context.WithTimeout(ctx, s.cfg.BlockInterval)
 	defer cancel()
-	maxTxs := s.cfg.maxTxsPerBlock()
-	var totalGas uint64
-	var txs [][]byte
-	for totalGas < s.cfg.MaxGasPerBlock && uint64(len(txs)) < maxTxs {
-		tx, err := utils.Recv(ctx, s.mempool)
-		if err != nil {
-			break
+
+	if s.txMempool.NumTxsNotPending() == 0 {
+		select {
+		case <-ctx.Done():
+		case <-s.txMempool.TxsAvailable():
 		}
-		txs = append(txs, tx.Payload)
-		totalGas += tx.GasUsed
+	}
+
+	txs, totalGas := s.txMempool.ReapMaxTxsBytesMaxGas(
+		int(s.cfg.maxTxsPerBlock()), // nolint:gosec // config values fit into int on supported platforms.
+		utils.Max[int64](),
+		int64(s.cfg.MaxGasPerBlock), // nolint:gosec // config values stay within int64 range.
+		int64(s.cfg.MaxGasPerBlock), // nolint:gosec // config values stay within int64 range.
+	)
+	s.txMempool.RemoveTxs(txs)
+	payloadTxs := make([][]byte, 0, len(txs))
+	for _, tx := range txs {
+		payloadTxs = append(payloadTxs, tx)
 	}
 	return types.PayloadBuilder{
 		CreatedAt: time.Now(),
-		TotalGas:  totalGas,
-		Txs:       txs,
+		// TODO: ReapMaxTxsBytesMaxGas does not handle corner cases correctly rn, which actually
+		// can produce negative total gas. Fixing it right away might be backward incompatible afaict,
+		// so we leave it as is for now.
+		TotalGas: uint64(totalGas), // nolint:gosec // guaranteed to be positive
+		Txs:      payloadTxs,
 	}.Build()
 }
 
@@ -86,11 +96,6 @@ func (s *State) nextPayload(ctx context.Context) (*types.Payload, error) {
 	}
 }
 
-// PushToMempool pushes the transaction to the mempool.
-func (s *State) PushToMempool(ctx context.Context, tx *pb.Transaction) error {
-	return utils.Send(ctx, s.mempool, tx)
-}
-
 // Run runs the background tasks of the producer state.
 func (s *State) Run(ctx context.Context) error {
 	return scope.Run(ctx, func(ctx context.Context, scope scope.Scope) error {

sei-tendermint/internal/consensus/mempool_test.go

@@ -20,6 +20,7 @@ import (
 	sm "github.com/sei-protocol/sei-chain/sei-tendermint/internal/state"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/store"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/test/factory"
+	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/types"
 )
 
@@ -247,7 +248,7 @@ func TestMempoolRmBadTx(t *testing.T) {
 
 		// check for the tx
 		for {
-			txs := cs.txMempool.ReapMaxBytesMaxGas(int64(len(txBytes)), -1, -1)
+			txs := cs.txMempool.ReapMaxBytesMaxGas(int64(len(txBytes)), utils.Max[int64](), utils.Max[int64]())
 			if len(txs) == 0 {
 				emptyMempoolCh <- struct{}{}
 				return

sei-tendermint/internal/mempool/mempool.go

@@ -18,13 +18,16 @@ import (
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils/scope"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/types"
+	"github.com/sei-protocol/seilog"
 )
 
-// errTxInCache is returned to the client if we saw tx earlier.
-var errTxInCache = errors.New("tx already exists in cache")
+var logger = seilog.NewLogger("tendermint", "internal", "mempool")
 
-// errTxTooLarge defines an error when a transaction is too big to be sent to peers.
-var errTxTooLarge = errors.New("tx too large")
+// ErrTxInCache is returned to the client if we saw tx earlier.
+var ErrTxInCache = errors.New("tx already exists in cache")
+
+// ErrTxTooLarge defines an error when a transaction is too big to be sent to peers.
+var ErrTxTooLarge = errors.New("tx too large")
 
 // Using SHA-256 truncated to 128 bits as the cache key: At 2K tx/sec, the
 // collision probability is effectively zero (≈10^-29 for 120K keys in a minute,
@@ -158,6 +161,10 @@ func NewTxMempool(
 	return txmp
 }
 
+func (txmp *TxMempool) Config() *config.MempoolConfig { return txmp.config }
+
+func (txmp *TxMempool) App() abci.Application { return txmp.app }
+
 func (txmp *TxMempool) TxStore() *TxStore { return txmp.txStore }
 
 // Lock obtains a write-lock on the mempool. A caller must be sure to explicitly
@@ -204,14 +211,14 @@ func (txmp *TxMempool) SizeBytes() int64 { return atomic.LoadInt64(&txmp.sizeByt
 
 func (txmp *TxMempool) PendingSizeBytes() int64 { return atomic.LoadInt64(&txmp.pendingSizeBytes) }
 
-// WaitForNextTx returns a blocking channel that will be closed when the next
-// valid transaction is available to gossip. It is thread-safe.
-func (txmp *TxMempool) WaitForNextTx() <-chan struct{} { return txmp.gossipIndex.WaitChan() }
-
-// NextGossipTx returns the next valid transaction to gossip. A caller must wait
-// for WaitForNextTx to signal a transaction is available to gossip first. It is
-// thread-safe.
-func (txmp *TxMempool) NextGossipTx() *clist.CElement { return txmp.gossipIndex.Front() }
+// WaitForNextTx waits until the next transaction is available for gossip.
+// Returns the next valid transaction to gossip.
+func (txmp *TxMempool) WaitForNextTx(ctx context.Context) (*clist.CElement, error) {
+	if _, _, err := utils.RecvOrClosed(ctx, txmp.gossipIndex.WaitChan()); err != nil {
+		return nil, err
+	}
+	return txmp.gossipIndex.Front(), nil
+}
 
 // TxsAvailable returns a channel which fires once for every height, and only
 // when transactions are available in the mempool. It is thread-safe.
@@ -269,14 +276,14 @@ func (txmp *TxMempool) CheckTx(
 	defer txmp.mtx.RUnlock()
 
 	if txSize := len(tx); txSize > txmp.config.MaxTxBytes {
-		return fmt.Errorf("%w: max size is %d, but got %d", errTxTooLarge, txmp.config.MaxTxBytes, txSize)
+		return fmt.Errorf("%w: max size is %d, but got %d", ErrTxTooLarge, txmp.config.MaxTxBytes, txSize)
 	}
 	constraints, err := txmp.txConstraintsFetcher()
 	if err != nil {
 		return fmt.Errorf("txmp.txConstraintsFetcher(): %w", err)
 	}
 	if txSize := types.ComputeProtoSizeForTxs([]types.Tx{tx}); txSize > constraints.MaxDataBytes {
-		return fmt.Errorf("%w: tx size is too big: %d, max: %d", errTxTooLarge, txSize, constraints.MaxDataBytes)
+		return fmt.Errorf("%w: tx size is too big: %d, max: %d", ErrTxTooLarge, txSize, constraints.MaxDataBytes)
 	}
 
 	// Reject low priority transactions when the mempool is more than
@@ -305,7 +312,7 @@ func (txmp *TxMempool) CheckTx(
 	// check if we've seen this transaction and error if we have.
 	if !txmp.cache.Push(txHash) {
 		txmp.txStore.GetOrSetPeerByTxHash(txHash, txInfo.SenderID)
-		return errTxInCache
+		return ErrTxInCache
 	}
 	txmp.metrics.CacheSize.Set(float64(txmp.cache.Size()))
 
@@ -482,7 +489,7 @@ func (txmp *TxMempool) Flush() {
 // and gas constraints. The returned list starts with EVM transactions (in priority order),
 // followed by non-EVM transactions (in priority order).
 // There are 4 types of constraints.
-//  1. maxBytes - stops pulling txs from mempool once maxBytes is hit. Can be set to -1 to be ignored.
+//  1. maxBytes - stops pulling txs from mempool once maxBytes is hit.
 //  2. maxGasWanted - stops pulling txs from mempool once total gas wanted exceeds maxGasWanted.
 //     Can be set to -1 to be ignored.
 //  3. maxGasEstimated - similar to maxGasWanted but will use the estimated gas used for EVM txs
@@ -492,6 +499,31 @@ func (txmp *TxMempool) Flush() {
 //   - Transactions returned are not removed from the mempool transaction
 //     store or indexes.
 func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGasWanted, maxGasEstimated int64) types.Txs {
+	txs, _ := txmp.ReapMaxTxsBytesMaxGas(utils.Max[int](), maxBytes, maxGasWanted, maxGasEstimated)
+	return txs
+}
+
+// ReapMaxTxsBytesMaxGas returns a list of transactions within the provided tx,
+// byte, and gas constraints together with the total estimated gas for the
+// returned transactions.
+//
+// NOTE: Gas limits are enforced using int64 running totals. If those totals
+// overflow, gas limit enforcement no longer works correctly. This preserves the
+// historical behavior for backward compatibility.
+func (txmp *TxMempool) ReapMaxTxsBytesMaxGas(maxTxs int, maxBytes, maxGasWanted, maxGasEstimated int64) (types.Txs, int64) {
+	if maxTxs < 0 {
+		maxTxs = utils.Max[int]()
+	}
+	if maxBytes < 0 {
+		maxBytes = utils.Max[int64]()
+	}
+	if maxGasWanted < 0 {
+		maxGasWanted = utils.Max[int64]()
+	}
+	if maxGasEstimated < 0 {
+		maxGasEstimated = utils.Max[int64]()
+	}
+
 	txmp.mtx.Lock()
 	defer txmp.mtx.Unlock()
 
@@ -505,7 +537,7 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGasWanted, maxGasEstimate
 	encounteredGasUnfit := false
 	if uint64(txmp.NumTxsNotPending()) < txmp.config.TxNotifyThreshold { //nolint:gosec // NumTxsNotPending returns non-negative value
 		// do not reap anything if threshold is not met
-		return []types.Tx{}
+		return []types.Tx{}, 0
 	}
 	totalTxs := txmp.priorityIndex.NumTxs()
 	evmTxs := make([]types.Tx, 0, totalTxs)
@@ -514,7 +546,7 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGasWanted, maxGasEstimate
 		size := types.ComputeProtoSizeForTxs([]types.Tx{wtx.tx})
 
 		// bytes limit is a hard stop
-		if maxBytes > -1 && totalSize+size > maxBytes {
+		if totalSize+size > maxBytes {
 			return false
 		}
 
@@ -531,8 +563,8 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGasWanted, maxGasEstimate
 		prospectiveGasWanted := totalGasWanted + wtx.gasWanted
 		prospectiveGasEstimated := totalGasEstimated + txGasEstimate
 
-		maxGasWantedExceeded := maxGasWanted > -1 && prospectiveGasWanted > maxGasWanted
-		maxGasEstimatedExceeded := maxGasEstimated > -1 && prospectiveGasEstimated > maxGasEstimated
+		maxGasWantedExceeded := prospectiveGasWanted > maxGasWanted
+		maxGasEstimatedExceeded := prospectiveGasEstimated > maxGasEstimated
 
 		if maxGasWantedExceeded || maxGasEstimatedExceeded {
 			// skip this unfit-by-gas tx once and attempt to pull up to 10 smaller ones
@@ -557,10 +589,25 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGasWanted, maxGasEstimate
 		if encounteredGasUnfit && numTxs >= MinTxsToPeek {
 			return false
 		}
+		if numTxs >= maxTxs {
+			return false
+		}
 		return true
 	})
 
-	return append(evmTxs, nonEvmTxs...)
+	return append(evmTxs, nonEvmTxs...), totalGasEstimated
+}
+
+// RemoveTxs removes the provided transactions from the mempool if present.
+func (txmp *TxMempool) RemoveTxs(txs types.Txs) {
+	txmp.Lock()
+	defer txmp.Unlock()
+
+	for _, tx := range txs {
+		if wtx := txmp.txStore.GetTxByHash(tx.Key()); wtx != nil {
+			txmp.removeTx(wtx, false, true, true)
+		}
+	}
 }
 
 // ReapMaxTxs returns a list of transactions within the provided number of