vms/platformvm: Prune mempool periodically (#2566)

Author: dhrubabasu

vms/platformvm/block/builder/builder.go

@@ -60,6 +60,13 @@ type Builder interface {
 
 	// BuildBlock can be called to attempt to create a new block
 	BuildBlock(context.Context) (snowman.Block, error)
+
+	// PackBlockTxs returns an array of txs that can fit into a valid block of
+	// size [targetBlockSize]. The returned txs are all verified against the
+	// preferred state.
+	//
+	// Note: This function does not call the consensus engine.
+	PackBlockTxs(targetBlockSize int) ([]*txs.Tx, error)
 }
 
 // builder implements a simple builder to convert txs into valid blocks
@@ -232,6 +239,24 @@ func (b *builder) BuildBlock(context.Context) (snowman.Block, error) {
 	return b.blkManager.NewBlock(statelessBlk), nil
 }
 
+func (b *builder) PackBlockTxs(targetBlockSize int) ([]*txs.Tx, error) {
+	preferredID := b.blkManager.Preferred()
+	preferredState, ok := b.blkManager.GetState(preferredID)
+	if !ok {
+		return nil, fmt.Errorf("%w: %s", errMissingPreferredState, preferredID)
+	}
+
+	return packBlockTxs(
+		preferredID,
+		preferredState,
+		b.Mempool,
+		b.txExecutorBackend,
+		b.blkManager,
+		b.txExecutorBackend.Clk.Time(),
+		targetBlockSize,
+	)
+}
+
 // [timestamp] is min(max(now, parent timestamp), next staker change time)
 func buildBlock(
 	builder *builder,
@@ -264,6 +289,7 @@ func buildBlock(
 				builder.txExecutorBackend,
 				builder.blkManager,
 				timestamp,
+				targetBlockSize,
 			)
 			if err != nil {
 				return nil, fmt.Errorf("failed to pack block txs: %w", err)
@@ -286,6 +312,7 @@ func buildBlock(
 		builder.txExecutorBackend,
 		builder.blkManager,
 		timestamp,
+		targetBlockSize,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to pack block txs: %w", err)
@@ -313,6 +340,7 @@ func packBlockTxs(
 	backend *txexecutor.Backend,
 	manager blockexecutor.Manager,
 	timestamp time.Time,
+	remainingSize int,
 ) ([]*txs.Tx, error) {
 	stateDiff, err := state.NewDiffOn(parentState)
 	if err != nil {
@@ -327,9 +355,8 @@ func packBlockTxs(
 	stateDiff.SetTimestamp(timestamp)
 
 	var (
-		blockTxs      []*txs.Tx
-		inputs        set.Set[ids.ID]
-		remainingSize = targetBlockSize
+		blockTxs []*txs.Tx
+		inputs   set.Set[ids.ID]
 	)
 
 	for {

vms/platformvm/config/execution_config.go

@@ -5,6 +5,7 @@ package config
 
 import (
 	"encoding/json"
+	"time"
 
 	"github.com/ava-labs/avalanchego/utils/units"
 	"github.com/ava-labs/avalanchego/vms/platformvm/network"
@@ -21,6 +22,7 @@ var DefaultExecutionConfig = ExecutionConfig{
 	BlockIDCacheSize:             8192,
 	FxOwnerCacheSize:             4 * units.MiB,
 	ChecksumsEnabled:             false,
+	MempoolPruneFrequency:        30 * time.Minute,
 }
 
 // ExecutionConfig provides execution parameters of PlatformVM
@@ -35,6 +37,7 @@ type ExecutionConfig struct {
 	BlockIDCacheSize             int            `json:"block-id-cache-size"`
 	FxOwnerCacheSize             int            `json:"fx-owner-cache-size"`
 	ChecksumsEnabled             bool           `json:"checksums-enabled"`
+	MempoolPruneFrequency        time.Duration  `json:"mempool-prune-frequency"`
 }
 
 // GetExecutionConfig returns an ExecutionConfig

vms/platformvm/config/execution_config_test.go

@@ -5,6 +5,7 @@ package config
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -61,7 +62,8 @@ func TestExecutionConfigUnmarshal(t *testing.T) {
 			"chain-db-cache-size": 7,
 			"block-id-cache-size": 8,
 			"fx-owner-cache-size": 9,
-			"checksums-enabled": true
+			"checksums-enabled": true,
+			"mempool-prune-frequency": 60000000000
 		}`)
 		ec, err := GetExecutionConfig(b)
 		require.NoError(err)
@@ -87,6 +89,7 @@ func TestExecutionConfigUnmarshal(t *testing.T) {
 			BlockIDCacheSize:             8,
 			FxOwnerCacheSize:             9,
 			ChecksumsEnabled:             true,
+			MempoolPruneFrequency:        time.Minute,
 		}
 		require.Equal(expected, ec)
 	})
@@ -135,6 +138,7 @@ func TestExecutionConfigUnmarshal(t *testing.T) {
 			BlockIDCacheSize:             8,
 			FxOwnerCacheSize:             9,
 			ChecksumsEnabled:             true,
+			MempoolPruneFrequency:        30 * time.Minute,
 		}
 		require.Equal(expected, ec)
 	})

vms/platformvm/vm.go

@@ -7,8 +7,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"net/http"
 	"sync"
+	"time"
 
 	"github.com/gorilla/rpc/v2"
 
@@ -247,6 +249,10 @@ func (vm *VM) Initialize(
 		return err
 	}
 
+	// Incrementing [awaitShutdown] would cause a deadlock since
+	// [periodicallyPruneMempool] grabs the context lock.
+	go vm.periodicallyPruneMempool(execConfig.MempoolPruneFrequency)
+
 	shouldPrune, err := vm.state.ShouldPrune()
 	if err != nil {
 		return fmt.Errorf(
@@ -274,6 +280,46 @@ func (vm *VM) Initialize(
 	return nil
 }
 
+func (vm *VM) periodicallyPruneMempool(frequency time.Duration) {
+	ticker := time.NewTicker(frequency)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-vm.onShutdownCtx.Done():
+			return
+		case <-ticker.C:
+			if err := vm.pruneMempool(); err != nil {
+				vm.ctx.Log.Debug("pruning mempool failed",
+					zap.Error(err),
+				)
+			}
+		}
+	}
+}
+
+func (vm *VM) pruneMempool() error {
+	vm.ctx.Lock.Lock()
+	defer vm.ctx.Lock.Unlock()
+
+	blockTxs, err := vm.Builder.PackBlockTxs(math.MaxInt)
+	if err != nil {
+		return err
+	}
+
+	for _, tx := range blockTxs {
+		if err := vm.Builder.Add(tx); err != nil {
+			vm.ctx.Log.Debug(
+				"failed to reissue tx",
+				zap.Stringer("txID", tx.ID()),
+				zap.Error(err),
+			)
+		}
+	}
+
+	return nil
+}
+
 // Create all chains that exist that this node validates.
 func (vm *VM) initBlockchains() error {
 	if vm.Config.PartialSyncPrimaryNetwork {

vms/platformvm/vm_test.go

@@ -2271,3 +2271,82 @@ func TestBaseTx(t *testing.T) {
 	require.NoError(baseTxBlock.Accept(context.Background()))
 	require.NoError(vm.SetPreference(context.Background(), vm.manager.LastAccepted()))
 }
+
+func TestPruneMempool(t *testing.T) {
+	require := require.New(t)
+	vm, _, _ := defaultVM(t, latestFork)
+	vm.ctx.Lock.Lock()
+	defer func() {
+		require.NoError(vm.Shutdown(context.Background()))
+		vm.ctx.Lock.Unlock()
+	}()
+
+	// Create a tx that will be valid regardless of timestamp.
+	sendAmt := uint64(100000)
+	changeAddr := ids.ShortEmpty
+
+	baseTx, err := vm.txBuilder.NewBaseTx(
+		sendAmt,
+		secp256k1fx.OutputOwners{
+			Threshold: 1,
+			Addrs: []ids.ShortID{
+				keys[1].Address(),
+			},
+		},
+		[]*secp256k1.PrivateKey{keys[0]},
+		changeAddr,
+	)
+	require.NoError(err)
+
+	vm.ctx.Lock.Unlock()
+	require.NoError(vm.issueTx(context.Background(), baseTx))
+	vm.ctx.Lock.Lock()
+
+	// [baseTx] should be in the mempool.
+	baseTxID := baseTx.ID()
+	_, ok := vm.Builder.Get(baseTxID)
+	require.True(ok)
+
+	// Create a tx that will be invalid after time advancement.
+	var (
+		startTime = vm.clock.Time()
+		endTime   = startTime.Add(vm.MinStakeDuration)
+	)
+
+	addValidatorTx, err := vm.txBuilder.NewAddValidatorTx(
+		defaultMinValidatorStake,
+		uint64(startTime.Unix()),
+		uint64(endTime.Unix()),
+		ids.GenerateTestNodeID(),
+		keys[2].Address(),
+		20000,
+		[]*secp256k1.PrivateKey{keys[1]},
+		ids.ShortEmpty,
+	)
+	require.NoError(err)
+
+	vm.ctx.Lock.Unlock()
+	require.NoError(vm.issueTx(context.Background(), addValidatorTx))
+	vm.ctx.Lock.Lock()
+
+	// Advance clock to [endTime], making [addValidatorTx] invalid.
+	vm.clock.Set(endTime)
+
+	// [addValidatorTx] and [baseTx] should still be in the mempool.
+	addValidatorTxID := addValidatorTx.ID()
+	_, ok = vm.Builder.Get(addValidatorTxID)
+	require.True(ok)
+	_, ok = vm.Builder.Get(baseTxID)
+	require.True(ok)
+
+	vm.ctx.Lock.Unlock()
+	require.NoError(vm.pruneMempool())
+	vm.ctx.Lock.Lock()
+
+	// [addValidatorTx] should be ejected from the mempool.
+	// [baseTx] should still be in the mempool.
+	_, ok = vm.Builder.Get(addValidatorTxID)
+	require.False(ok)
+	_, ok = vm.Builder.Get(baseTxID)
+	require.True(ok)
+}