vms/avm: Simplify Peek function in mempool (#2449)

Co-authored-by: Stephen Buttolph <stephen@avalabs.org>

Author: dhrubabasu

vms/avm/block/builder/builder.go

@@ -93,8 +93,12 @@ func (b *builder) BuildBlock(context.Context) (snowman.Block, error) {
 		remainingSize = targetBlockSize
 	)
 	for {
-		tx := b.mempool.Peek(remainingSize)
-		if tx == nil {
+		tx := b.mempool.Peek()
+		// Invariant: [mempool.MaxTxSize] < [targetBlockSize]. This guarantees
+		// that we will only stop building a block once there are no
+		// transactions in the mempool or the block is at least
+		// [targetBlockSize - mempool.MaxTxSize] bytes full.
+		if tx == nil || len(tx.Bytes()) > remainingSize {
 			break
 		}
 		b.mempool.Remove([]*txs.Tx{tx})

vms/avm/block/builder/builder_test.go

@@ -134,11 +134,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				tx := &txs.Tx{Unsigned: unsignedTx}
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(gomock.Any()).Return(tx)
+				mempool.EXPECT().Peek().Return(tx)
 				mempool.EXPECT().Remove([]*txs.Tx{tx})
 				mempool.EXPECT().MarkDropped(tx.ID(), errTest)
 				// Second loop iteration
-				mempool.EXPECT().Peek(gomock.Any()).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				return New(
@@ -179,11 +179,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				tx := &txs.Tx{Unsigned: unsignedTx}
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(gomock.Any()).Return(tx)
+				mempool.EXPECT().Peek().Return(tx)
 				mempool.EXPECT().Remove([]*txs.Tx{tx})
 				mempool.EXPECT().MarkDropped(tx.ID(), errTest)
 				// Second loop iteration
-				mempool.EXPECT().Peek(gomock.Any()).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				return New(
@@ -225,11 +225,11 @@ func TestBuilderBuildBlock(t *testing.T) {
 				tx := &txs.Tx{Unsigned: unsignedTx}
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(gomock.Any()).Return(tx)
+				mempool.EXPECT().Peek().Return(tx)
 				mempool.EXPECT().Remove([]*txs.Tx{tx})
 				mempool.EXPECT().MarkDropped(tx.ID(), errTest)
 				// Second loop iteration
-				mempool.EXPECT().Peek(gomock.Any()).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				return New(
@@ -309,14 +309,14 @@ func TestBuilderBuildBlock(t *testing.T) {
 				)
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(targetBlockSize).Return(tx1)
+				mempool.EXPECT().Peek().Return(tx1)
 				mempool.EXPECT().Remove([]*txs.Tx{tx1})
 				// Second loop iteration
-				mempool.EXPECT().Peek(targetBlockSize - len(tx1Bytes)).Return(tx2)
+				mempool.EXPECT().Peek().Return(tx2)
 				mempool.EXPECT().Remove([]*txs.Tx{tx2})
 				mempool.EXPECT().MarkDropped(tx2.ID(), blkexecutor.ErrConflictingBlockTxs)
 				// Third loop iteration
-				mempool.EXPECT().Peek(targetBlockSize - len(tx1Bytes)).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				// To marshal the tx/block
@@ -385,10 +385,10 @@ func TestBuilderBuildBlock(t *testing.T) {
 				tx := &txs.Tx{Unsigned: unsignedTx}
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(gomock.Any()).Return(tx)
+				mempool.EXPECT().Peek().Return(tx)
 				mempool.EXPECT().Remove([]*txs.Tx{tx})
 				// Second loop iteration
-				mempool.EXPECT().Peek(gomock.Any()).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				// To marshal the tx/block
@@ -459,10 +459,10 @@ func TestBuilderBuildBlock(t *testing.T) {
 				tx := &txs.Tx{Unsigned: unsignedTx}
 
 				mempool := mempool.NewMockMempool(ctrl)
-				mempool.EXPECT().Peek(gomock.Any()).Return(tx)
+				mempool.EXPECT().Peek().Return(tx)
 				mempool.EXPECT().Remove([]*txs.Tx{tx})
 				// Second loop iteration
-				mempool.EXPECT().Peek(gomock.Any()).Return(nil)
+				mempool.EXPECT().Peek().Return(nil)
 				mempool.EXPECT().RequestBuildBlock()
 
 				// To marshal the tx/block

vms/avm/txs/mempool/mempool.go

@@ -48,8 +48,8 @@ type Mempool interface {
 	Get(txID ids.ID) *txs.Tx
 	Remove(txs []*txs.Tx)
 
-	// Peek returns the first tx in the mempool whose size is <= [maxTxSize].
-	Peek(maxTxSize int) *txs.Tx
+	// Peek returns the oldest tx in the mempool.
+	Peek() *txs.Tx
 
 	// RequestBuildBlock notifies the consensus engine that a block should be
 	// built if there is at least one transaction in the mempool.
@@ -182,16 +182,9 @@ func (m *mempool) Remove(txsToRemove []*txs.Tx) {
 	}
 }
 
-func (m *mempool) Peek(maxTxSize int) *txs.Tx {
-	txIter := m.unissuedTxs.NewIterator()
-	for txIter.Next() {
-		tx := txIter.Value()
-		txSize := len(tx.Bytes())
-		if txSize <= maxTxSize {
-			return tx
-		}
-	}
-	return nil
+func (m *mempool) Peek() *txs.Tx {
+	_, tx, _ := m.unissuedTxs.Oldest()
+	return tx
 }
 
 func (m *mempool) RequestBuildBlock() {

vms/avm/txs/mempool/mempool_test.go

@@ -170,3 +170,31 @@ func createTestTxs(count int) []*txs.Tx {
 	}
 	return testTxs
 }
+
+func TestPeekTxs(t *testing.T) {
+	require := require.New(t)
+
+	registerer := prometheus.NewRegistry()
+	toEngine := make(chan common.Message, 100)
+	mempool, err := New("mempool", registerer, toEngine)
+	require.NoError(err)
+
+	testTxs := createTestTxs(2)
+
+	require.Nil(mempool.Peek())
+
+	require.NoError(mempool.Add(testTxs[0]))
+	require.NoError(mempool.Add(testTxs[1]))
+
+	require.Equal(mempool.Peek(), testTxs[0])
+	require.NotEqual(mempool.Peek(), testTxs[1])
+
+	mempool.Remove([]*txs.Tx{testTxs[0]})
+
+	require.NotEqual(mempool.Peek(), testTxs[0])
+	require.Equal(mempool.Peek(), testTxs[1])
+
+	mempool.Remove([]*txs.Tx{testTxs[1]})
+
+	require.Nil(mempool.Peek())
+}