Add SetPreferenceWithContext

snow/engine/snowman/block/block_context_vm.go

@@ -8,6 +8,7 @@ package block
 import (
 	"context"
 
+	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
 )
 
@@ -39,6 +40,12 @@ type BuildBlockWithContextChainVM interface {
 	BuildBlockWithContext(ctx context.Context, blockCtx *Context) (snowman.Block, error)
 }
 
+// SetPreferenceWithContextChainVM defines the interface a ChainVM can optionally
+// implement to consider the P-Chain height when setting preference.
+type SetPreferenceWithContextChainVM interface {
+	SetPreferenceWithContext(ctx context.Context, blkID ids.ID, blockCtx *Context) error
+}
+
 // WithVerifyContext defines the interface a Block can optionally implement to
 // consider the P-Chain height when verifying itself.
 //

snow/engine/snowman/block/blocktest/set_preference_vm.go

@@ -0,0 +1,42 @@
+// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package blocktest
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
+)
+
+var (
+	errSetPreferenceWithContext = errors.New("unexpectedly called SetPreferenceWithContext")
+
+	_ block.SetPreferenceWithContextChainVM = (*SetPreferenceVM)(nil)
+)
+
+type SetPreferenceVM struct {
+	T *testing.T
+
+	CantSetPreferenceWithContext bool
+	SetPreferenceWithContextF    func(context.Context, ids.ID, *block.Context) error
+}
+
+func (vm *SetPreferenceVM) Default(cant bool) {
+	vm.CantSetPreferenceWithContext = cant
+}
+
+func (vm *SetPreferenceVM) SetPreferenceWithContext(ctx context.Context, id ids.ID, blockCtx *block.Context) error {
+	if vm.SetPreferenceWithContextF != nil {
+		return vm.SetPreferenceWithContextF(ctx, id, blockCtx)
+	}
+	if vm.CantSetPreferenceWithContext && vm.T != nil {
+		require.FailNow(vm.T, errSetPreferenceWithContext.Error())
+	}
+	return errSetPreferenceWithContext
+}

vms/metervm/block_metrics.go

@@ -32,6 +32,8 @@ type blockMetrics struct {
 	// Block building with context metrics
 	buildBlockWithContext,
 	buildBlockWithContextErr,
+	// Setting preference with context metrics
+	setPreferenceWithContext,
 	// Batched metrics
 	getAncestors,
 	batchedParseBlock,
@@ -47,6 +49,7 @@ type blockMetrics struct {
 
 func (m *blockMetrics) Initialize(
 	supportsBlockBuildingWithContext bool,
+	supportsSettingPreferenceWithContext bool,
 	supportsBatchedFetching bool,
 	supportsStateSync bool,
 	reg prometheus.Registerer,
@@ -73,6 +76,9 @@ func (m *blockMetrics) Initialize(
 		m.buildBlockWithContext = newAverager("build_block_with_context", reg, &errs)
 		m.buildBlockWithContextErr = newAverager("build_block_with_context_err", reg, &errs)
 	}
+	if supportsSettingPreferenceWithContext {
+		m.setPreferenceWithContext = newAverager("set_preference_with_context", reg, &errs)
+	}
 	if supportsBatchedFetching {
 		m.getAncestors = newAverager("get_ancestors", reg, &errs)
 		m.batchedParseBlock = newAverager("batched_parse_block", reg, &errs)

vms/metervm/block_vm.go

@@ -18,17 +18,19 @@ import (
 )
 
 var (
-	_ block.ChainVM                      = (*blockVM)(nil)
-	_ block.BuildBlockWithContextChainVM = (*blockVM)(nil)
-	_ block.BatchedChainVM               = (*blockVM)(nil)
-	_ block.StateSyncableVM              = (*blockVM)(nil)
+	_ block.ChainVM                         = (*blockVM)(nil)
+	_ block.BuildBlockWithContextChainVM    = (*blockVM)(nil)
+	_ block.SetPreferenceWithContextChainVM = (*blockVM)(nil)
+	_ block.BatchedChainVM                  = (*blockVM)(nil)
+	_ block.StateSyncableVM                 = (*blockVM)(nil)
 )
 
 type blockVM struct {
 	block.ChainVM
-	buildBlockVM block.BuildBlockWithContextChainVM
-	batchedVM    block.BatchedChainVM
-	ssVM         block.StateSyncableVM
+	buildBlockVM    block.BuildBlockWithContextChainVM
+	setPreferenceVM block.SetPreferenceWithContextChainVM
+	batchedVM       block.BatchedChainVM
+	ssVM            block.StateSyncableVM
 
 	blockMetrics
 	registry prometheus.Registerer
@@ -39,14 +41,16 @@ func NewBlockVM(
 	reg prometheus.Registerer,
 ) block.ChainVM {
 	buildBlockVM, _ := vm.(block.BuildBlockWithContextChainVM)
+	setPreferenceVM, _ := vm.(block.SetPreferenceWithContextChainVM)
 	batchedVM, _ := vm.(block.BatchedChainVM)
 	ssVM, _ := vm.(block.StateSyncableVM)
 	return &blockVM{
-		ChainVM:      vm,
-		buildBlockVM: buildBlockVM,
-		batchedVM:    batchedVM,
-		ssVM:         ssVM,
-		registry:     reg,
+		ChainVM:         vm,
+		buildBlockVM:    buildBlockVM,
+		setPreferenceVM: setPreferenceVM,
+		batchedVM:       batchedVM,
+		ssVM:            ssVM,
+		registry:        reg,
 	}
 }
 
@@ -62,6 +66,7 @@ func (vm *blockVM) Initialize(
 ) error {
 	err := vm.blockMetrics.Initialize(
 		vm.buildBlockVM != nil,
+		vm.setPreferenceVM != nil,
 		vm.batchedVM != nil,
 		vm.ssVM != nil,
 		vm.registry,

vms/metervm/set_preference_with_context_vm.go

@@ -0,0 +1,23 @@
+// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package metervm
+
+import (
+	"context"
+	"time"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/ava-labs/avalanchego/snow/engine/snowman/block"
+)
+
+func (vm *blockVM) SetPreferenceWithContext(ctx context.Context, blkID ids.ID, blockCtx *block.Context) error {
+	if vm.setPreferenceVM == nil {
+		return vm.SetPreference(ctx, blkID)
+	}
+
+	start := time.Now()
+	err := vm.setPreferenceVM.SetPreferenceWithContext(ctx, blkID, blockCtx)
+	vm.blockMetrics.setPreferenceWithContext.Observe(float64(time.Since(start)))
+	return err
+}

vms/platformvm/block/builder/builder_test.go

@@ -157,7 +157,7 @@ func TestBuildBlockShouldReward(t *testing.T) {
 	require.Equal([]*txs.Tx{tx}, blk.(*blockexecutor.Block).Block.Txs())
 	require.NoError(blk.Verify(context.Background()))
 	require.NoError(blk.Accept(context.Background()))
-	env.blkManager.SetPreference(blk.ID())
+	env.blkManager.SetPreference(blk.ID(), nil)
 
 	// Validator should now be current
 	staker, err := env.state.GetCurrentValidator(constants.PrimaryNetworkID, nodeID)
@@ -196,7 +196,7 @@ func TestBuildBlockShouldReward(t *testing.T) {
 		require.NoError(blk.Accept(context.Background()))
 		require.NoError(commit.Verify(context.Background()))
 		require.NoError(commit.Accept(context.Background()))
-		env.blkManager.SetPreference(commit.ID())
+		env.blkManager.SetPreference(commit.ID(), nil)
 
 		// Stop rewarding once our staker is rewarded
 		if staker.TxID == txID {
@@ -459,7 +459,7 @@ func TestNoErrorOnUnexpectedSetPreferenceDuringBootstrapping(t *testing.T) {
 	defer env.ctx.Lock.Unlock()
 
 	env.isBootstrapped.Set(false)
-	env.blkManager.SetPreference(ids.GenerateTestID()) // should not panic
+	env.blkManager.SetPreference(ids.GenerateTestID(), nil) // should not panic
 }
 
 func TestGetNextStakerToReward(t *testing.T) {

vms/platformvm/block/builder/helpers_test.go

@@ -178,7 +178,7 @@ func newEnvironment(t *testing.T, f upgradetest.Fork) *environment { //nolint:un
 		res.blkManager,
 	)
 
-	res.blkManager.SetPreference(genesisID)
+	res.blkManager.SetPreference(genesisID, nil)
 	addSubnet(t, res)
 
 	t.Cleanup(func() {

vms/platformvm/block/executor/manager.go

@@ -19,6 +19,8 @@ import (
 	"github.com/ava-labs/avalanchego/vms/platformvm/txs/fee"
 	"github.com/ava-labs/avalanchego/vms/platformvm/validators"
 	"github.com/ava-labs/avalanchego/vms/txs/mempool"
+
+	snowmanblock "github.com/ava-labs/avalanchego/snow/engine/snowman/block"
 )
 
 var (
@@ -34,7 +36,7 @@ type Manager interface {
 	// Returns the ID of the most recently accepted block.
 	LastAccepted() ids.ID
 
-	SetPreference(blkID ids.ID)
+	SetPreference(blkID ids.ID, blockCtx *snowmanblock.Context)
 	Preferred() ids.ID
 
 	GetBlock(blkID ids.ID) (snowman.Block, error)
@@ -88,6 +90,7 @@ type manager struct {
 	rejector block.Visitor
 
 	preferred         ids.ID
+	preferredCtx      *snowmanblock.Context
 	txExecutorBackend *executor.Backend
 }
 
@@ -110,8 +113,9 @@ func (m *manager) NewBlock(blk block.Block) snowman.Block {
 	}
 }
 
-func (m *manager) SetPreference(blkID ids.ID) {
+func (m *manager) SetPreference(blkID ids.ID, blockCtx *snowmanblock.Context) {
 	m.preferred = blkID
+	m.preferredCtx = blockCtx
 }
 
 func (m *manager) Preferred() ids.ID {
@@ -132,9 +136,20 @@ func (m *manager) VerifyTx(tx *txs.Tx) error {
 		}
 	}
 
-	recommendedPChainHeight, err := m.ctx.ValidatorState.GetMinimumHeight(context.TODO())
-	if err != nil {
-		return fmt.Errorf("failed to fetch P-chain height: %w", err)
+	var (
+		recommendedPChainHeight uint64
+		err                     error
+	)
+	if m.preferredCtx != nil {
+		recommendedPChainHeight = m.preferredCtx.PChainHeight
+	} else {
+		// TODO: After Granite activation, what P-Chain height should be used
+		// if no preference has been set yet? This is the case immediately following
+		// initialization before any SetPreferenceWithContext calls have been made.
+		recommendedPChainHeight, err = m.ctx.ValidatorState.GetMinimumHeight(context.TODO())
+		if err != nil {
+			return fmt.Errorf("failed to fetch P-chain height: %w", err)
+		}
 	}
 	err = executor.VerifyWarpMessages(
 		context.TODO(),

vms/platformvm/block/executor/manager_test.go

@@ -13,6 +13,8 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/vms/platformvm/block"
 	"github.com/ava-labs/avalanchego/vms/platformvm/state"
+
+	snowmanblock "github.com/ava-labs/avalanchego/snow/engine/snowman/block"
 )
 
 func TestGetBlock(t *testing.T) {
@@ -82,6 +84,25 @@ func TestManagerSetPreference(t *testing.T) {
 	require.Equal(initialPreference, manager.Preferred())
 
 	newPreference := ids.GenerateTestID()
-	manager.SetPreference(newPreference)
+	manager.SetPreference(newPreference, nil)
+	require.Equal(newPreference, manager.Preferred())
+}
+
+func TestManagerSetPreferenceWithContext(t *testing.T) {
+	require := require.New(t)
+
+	initialPreference := ids.GenerateTestID()
+	manager := &manager{
+		preferred: initialPreference,
+	}
+	require.Equal(initialPreference, manager.Preferred())
+	require.Nil(manager.preferredCtx)
+
+	newPreference := ids.GenerateTestID()
+	newContext := &snowmanblock.Context{
+		PChainHeight: 100,
+	}
+	manager.SetPreference(newPreference, newContext)
 	require.Equal(newPreference, manager.Preferred())
+	require.Equal(newContext, manager.preferredCtx)
 }

vms/platformvm/block/executor/proposal_block_test.go

@@ -1372,7 +1372,7 @@ func TestAddValidatorProposalBlock(t *testing.T) {
 	blk := env.blkManager.NewBlock(statelessBlk)
 	require.NoError(blk.Verify(context.Background()))
 	require.NoError(blk.Accept(context.Background()))
-	env.blkManager.SetPreference(statelessBlk.ID())
+	env.blkManager.SetPreference(statelessBlk.ID(), nil)
 
 	// Should be current
 	staker, err := env.state.GetCurrentValidator(constants.PrimaryNetworkID, nodeID)
@@ -1405,7 +1405,7 @@ func TestAddValidatorProposalBlock(t *testing.T) {
 		blk = env.blkManager.NewBlock(statelessBlk)
 		require.NoError(blk.Verify(context.Background()))
 		require.NoError(blk.Accept(context.Background()))
-		env.blkManager.SetPreference(statelessBlk.ID())
+		env.blkManager.SetPreference(statelessBlk.ID(), nil)
 	}
 
 	env.clk.Set(validatorEndTime)