Update ReapMaxBytesMaxGas to include estimated gas (#263)

* checkpoint: add gasEstimate to WrappedTx in mempool

* updated proto with gas estimated + build blocks with gas used + mempool tests

* fix replay stub

* implement fallback to gasWanted if gasEstimate is not set

* remove print

* fix some tests

* fix some tests

* fix nondeterministic test

* remove both maxGasWanted and maxGasEstimated, just use maxGas

* fix test: mempool mocks

* fix race: pending tx mempool size bytes

Author: jewei1997

abci/types/types.pb.go

@@ -198,7 +198,7 @@ func (txmp *TxMempool) BytesNotPending() int64 {
 }
 
 func (txmp *TxMempool) TotalTxsBytesSize() int64 {
-	return txmp.BytesNotPending() + int64(txmp.pendingTxs.sizeBytes)
+	return txmp.BytesNotPending() + int64(txmp.pendingTxs.SizeBytes())
 }
 
 // PendingSize returns the number of pending transactions in the mempool.
@@ -335,6 +335,7 @@ func (txmp *TxMempool) CheckTx(
 		evmAddress:    res.EVMSenderAddress,
 		isEVM:         res.IsEVM,
 		removeHandler: removeHandler,
+		estimatedGas:  res.GasEstimated,
 	}
 
 	if err == nil {
@@ -452,9 +453,19 @@ func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGas, minTxsInBlock int64)
 		if maxBytes > -1 && totalSize+size > maxBytes {
 			return false
 		}
+
+		// if the tx doesn't have a gas estimate, fallback to gas wanted
+		var gasEstimateOrWanted int64
+		if wtx.estimatedGas > 0 {
+			gasEstimateOrWanted = wtx.estimatedGas
+		} else {
+			gasEstimateOrWanted = wtx.gasWanted
+		}
+
 		totalSize += size
-		gas := totalGas + wtx.gasWanted
-		if nonzeroGasTxCnt >= minTxsInBlock && maxGas > -1 && gas > maxGas {
+		gas := totalGas + gasEstimateOrWanted
+		maxGasExceeded := maxGas > -1 && gas > maxGas
+		if nonzeroGasTxCnt >= minTxsInBlock && maxGasExceeded {
 			return false
 		}
 
@@ -679,6 +690,7 @@ func (txmp *TxMempool) addNewTransaction(wtx *WrappedTx, res *abci.ResponseCheck
 	}
 
 	wtx.gasWanted = res.GasWanted
+	wtx.estimatedGas = res.GasEstimated
 	wtx.priority = priority
 	wtx.sender = sender
 	wtx.peers = map[uint16]struct{}{

internal/mempool/mempool.go

@@ -30,6 +30,8 @@ import (
 type application struct {
 	*kvstore.Application
 
+	gasWanted      *int64
+	gasEstimated   *int64
 	occupiedNonces map[string][]uint64
 }
 
@@ -38,13 +40,26 @@ type testTx struct {
 	priority int64
 }
 
+var DefaultGasEstimated = int64(1)
+var DefaultGasWanted = int64(1)
+
 func (app *application) CheckTx(_ context.Context, req *abci.RequestCheckTx) (*abci.ResponseCheckTxV2, error) {
 
 	var (
 		priority int64
 		sender   string
 	)
 
+	gasWanted := DefaultGasWanted
+	if app.gasWanted != nil {
+		gasWanted = *app.gasWanted
+	}
+
+	gasEstimated := DefaultGasEstimated
+	if app.gasEstimated != nil {
+		gasEstimated = *app.gasEstimated
+	}
+
 	if strings.HasPrefix(string(req.Tx), "evm") {
 		// format is evm-sender-0=account=priority=nonce
 		// split into respective vars
@@ -55,18 +70,20 @@ func (app *application) CheckTx(_ context.Context, req *abci.RequestCheckTx) (*a
 		if err != nil {
 			// could not parse
 			return &abci.ResponseCheckTxV2{ResponseCheckTx: &abci.ResponseCheckTx{
-				Priority:  priority,
-				Code:      100,
-				GasWanted: 1,
+				Priority:     priority,
+				Code:         100,
+				GasWanted:    gasWanted,
+				GasEstimated: gasEstimated,
 			}}, nil
 		}
 		nonce, err := strconv.ParseUint(string(parts[3]), 10, 64)
 		if err != nil {
 			// could not parse
 			return &abci.ResponseCheckTxV2{ResponseCheckTx: &abci.ResponseCheckTx{
-				Priority:  priority,
-				Code:      101,
-				GasWanted: 1,
+				Priority:     priority,
+				Code:         101,
+				GasWanted:    gasWanted,
+				GasEstimated: gasEstimated,
 			}}, nil
 		}
 		if app.occupiedNonces == nil {
@@ -92,9 +109,10 @@ func (app *application) CheckTx(_ context.Context, req *abci.RequestCheckTx) (*a
 		app.occupiedNonces[account] = append(app.occupiedNonces[account], nonce)
 		return &abci.ResponseCheckTxV2{
 			ResponseCheckTx: &abci.ResponseCheckTx{
-				Priority:  v,
-				Code:      code.CodeTypeOK,
-				GasWanted: 1,
+				Priority:     v,
+				Code:         code.CodeTypeOK,
+				GasWanted:    gasWanted,
+				GasEstimated: gasEstimated,
 			},
 			EVMNonce:             nonce,
 			EVMSenderAddress:     account,
@@ -122,26 +140,29 @@ func (app *application) CheckTx(_ context.Context, req *abci.RequestCheckTx) (*a
 		v, err := strconv.ParseInt(string(parts[2]), 10, 64)
 		if err != nil {
 			return &abci.ResponseCheckTxV2{ResponseCheckTx: &abci.ResponseCheckTx{
-				Priority:  priority,
-				Code:      100,
-				GasWanted: 1,
+				Priority:     priority,
+				Code:         100,
+				GasWanted:    gasWanted,
+				GasEstimated: gasEstimated,
 			}}, nil
 		}
 
 		priority = v
 		sender = string(parts[0])
 	} else {
 		return &abci.ResponseCheckTxV2{ResponseCheckTx: &abci.ResponseCheckTx{
-			Priority:  priority,
-			Code:      101,
-			GasWanted: 1,
+			Priority:     priority,
+			Code:         101,
+			GasWanted:    gasWanted,
+			GasEstimated: gasEstimated,
 		}}, nil
 	}
 	return &abci.ResponseCheckTxV2{ResponseCheckTx: &abci.ResponseCheckTx{
-		Priority:  priority,
-		Sender:    sender,
-		Code:      code.CodeTypeOK,
-		GasWanted: 1,
+		Priority:     priority,
+		Sender:       sender,
+		Code:         code.CodeTypeOK,
+		GasWanted:    gasWanted,
+		GasEstimated: gasEstimated,
 	}}, nil
 }
 
@@ -346,7 +367,9 @@ func TestTxMempool_ReapMaxBytesMaxGas(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	client := abciclient.NewLocalClient(log.NewNopLogger(), &application{Application: kvstore.NewApplication()})
+	// gas estimated is 1 so we will use this and not fallback to gas wanted
+	gasEstimated := int64(1)
+	client := abciclient.NewLocalClient(log.NewNopLogger(), &application{Application: kvstore.NewApplication(), gasEstimated: &gasEstimated})
 	if err := client.Start(ctx); err != nil {
 		t.Fatal(err)
 	}
@@ -424,6 +447,66 @@ func TestTxMempool_ReapMaxBytesMaxGas(t *testing.T) {
 		require.Len(t, reapedTxs, 2)
 	}()
 
+	// Reap by max gas estimated
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxBytesMaxGas(-1, 50, 0)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Len(t, reapedTxs, 50)
+	}()
+
+	wg.Wait()
+}
+
+func TestTxMempool_ReapMaxBytesMaxGas_FallbackToGasWanted(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// gas estimated is 0 so we will fallback to gas wanted
+	gasEstimated := int64(0)
+	client := abciclient.NewLocalClient(log.NewNopLogger(), &application{Application: kvstore.NewApplication(), gasEstimated: &gasEstimated})
+	if err := client.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(client.Wait)
+
+	txmp := setup(t, client, 0)
+	tTxs := checkTxs(ctx, t, txmp, 100, 0)
+
+	txMap := make(map[types.TxKey]testTx)
+	priorities := make([]int64, len(tTxs))
+	for i, tTx := range tTxs {
+		txMap[tTx.tx.Key()] = tTx
+		priorities[i] = tTx.priority
+	}
+
+	// Debug: Print sorted priorities
+	sort.Slice(priorities, func(i, j int) bool {
+		return priorities[i] > priorities[j]
+	})
+
+	ensurePrioritized := func(reapedTxs types.Txs) {
+		reapedPriorities := make([]int64, len(reapedTxs))
+		for i, rTx := range reapedTxs {
+			reapedPriorities[i] = txMap[rTx.Key()].priority
+		}
+
+		require.Equal(t, priorities[:len(reapedPriorities)], reapedPriorities)
+	}
+
+	var wg sync.WaitGroup
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		reapedTxs := txmp.ReapMaxBytesMaxGas(-1, 50, 0)
+		ensurePrioritized(reapedTxs)
+		require.Equal(t, len(tTxs), txmp.Size())
+		require.Len(t, reapedTxs, 50)
+	}()
+
 	wg.Wait()
 }
 

internal/mempool/mempool_test.go

@@ -39,6 +39,9 @@ type WrappedTx struct {
 	// gasWanted defines the amount of gas the transaction sender requires
 	gasWanted int64
 
+	// estimatedGas defines the amount of gas that the transaction is estimated to use
+	estimatedGas int64
+
 	// priority defines the transaction's priority as specified by the application
 	// in the ResponseCheckTx response.
 	priority int64
@@ -397,6 +400,12 @@ func (p *PendingTxs) Insert(tx *WrappedTx, resCheckTx *abci.ResponseCheckTxV2, t
 	return nil
 }
 
+func (p *PendingTxs) SizeBytes() uint64 {
+	p.mtx.RLock()
+	defer p.mtx.RUnlock()
+	return p.sizeBytes
+}
+
 func (p *PendingTxs) Peek(max int) []TxWithResponse {
 	p.mtx.RLock()
 	defer p.mtx.RUnlock()

internal/mempool/mocks/mempool.go

@@ -44,9 +44,9 @@ type Mempool interface {
 
 	// ReapMaxBytesMaxGas reaps transactions from the mempool up to maxBytes
 	// bytes total with the condition that the total gasWanted must be less than
-	// maxGas.
+	// maxGas and that the total estimated gas used is less than maxGasEstimated.
 	//
-	// If both maxes are negative, there is no cap on the size of all returned
+	// If all 3 maxes are negative, there is no cap on the size of all returned
 	// transactions (~ all available transactions).
 	ReapMaxBytesMaxGas(maxBytes, maxGas, minTxsInBlock int64) types.Txs
 

internal/mempool/tx.go

@@ -994,8 +994,10 @@ func TestPeerManager_Dialed_MaxConnectedUpgrade(t *testing.T) {
 	// Start upgrade with c
 	dial, err := peerManager.TryDialNext()
 	require.NoError(t, err)
-	require.Equal(t, c, dial)
-	require.NoError(t, peerManager.Dialed(c))
+	if dial != c && dial != d {
+		t.Fatalf("dial = %s, expected %s or %s", dial, c, d)
+	}
+	require.NoError(t, peerManager.Dialed(dial))
 
 	// Try to dial d - should fail since we're at upgrade capacity
 	dial, err = peerManager.TryDialNext()

internal/mempool/types.go

@@ -297,6 +297,7 @@ message ResponseCheckTx {
   string         codespace  = 8;
   string         sender     = 9;
   int64          priority   = 10;
+  int64          gas_estimated = 12;
 
   reserved 4, 6, 7, 11; // see https://github.com/tendermint/tendermint/issues/8543
 }