Build block immediately when transitioning epoch

adapters.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/ava-labs/simplex/common"
 	metadata "github.com/ava-labs/simplex/msm"
+	"github.com/ava-labs/simplex/simplex"
 )
 
 type Communication struct {
@@ -200,6 +201,7 @@ type BlockBuilderWaiter struct {
 	lock   sync.Mutex
 	cancel context.CancelFunc
 	msm    *metadata.StateMachine
+	e      *simplex.Epoch
 	vm     VM
 }
 
@@ -221,7 +223,9 @@ func (bw *BlockBuilderWaiter) WaitForPendingBlock(ctx context.Context) {
 	bw.cancel = cancel
 	bw.lock.Unlock()
 	defer cancel()
-	bw.vm.WaitForPendingBlock(ctx)
+
+	md := bw.e.Metadata()
+	bw.msm.WaitForPendingBlock(ctx, md)
 }
 
 func (bw *BlockBuilderWaiter) BuildBlock(ctx context.Context, metadata common.ProtocolMetadata, blacklist common.Blacklist) (common.VerifiedBlock, bool) {

instance.go

@@ -382,15 +382,15 @@ func (i *Instance) processEpochChange(epochChange epochChange) {
 
 // startEpoch starts a new epoch with the given configuration.
 // Must be called under the lock, and assumes that the previous epoch has been stopped (if any).
-func (i *Instance) startEpoch(epochConfig simplex.EpochConfig) error {
-	epoch, err := simplex.NewEpoch(epochConfig)
+func (i *Instance) startEpoch(epochConfig *epochConfig) error {
+	epoch, err := simplex.NewEpoch(epochConfig.EpochConfig)
 	if err != nil {
 		return fmt.Errorf("error creating simplex epoch: %w", err)
 	}
 	epoch.Epoch = epochConfig.Epoch
 	i.e = epoch
 	i.epochOrNV = epoch
-
+	epochConfig.bbw.e = epoch
 	return epoch.Start()
 }
 
@@ -417,30 +417,30 @@ func (i *Instance) iAmValidator(nodes common.Nodes) bool {
 	return false
 }
 
-func (i *Instance) createEpochConfig() (simplex.EpochConfig, error) {
+func (i *Instance) createEpochConfig() (*epochConfig, error) {
 	lastBlock, numBlocks, err := i.lastBlock()
 	if err != nil {
-		return simplex.EpochConfig{}, err
+		return nil, err
 	}
 
 	lastNonSimplexHeight := i.Config.LastNonSimplexInnerBlock.Height()
 	genesisValidatorSet := i.Config.PlatformChain.GenesisValidatorSet()
 	nodes, epochNum, err := constructEpochAndValidatorSet(i.Config.Logger, lastNonSimplexHeight, genesisValidatorSet, numBlocks, &ParsedBlock{StateMachineBlock: lastBlock}, i.Config.Storage)
 	if err != nil {
-		return simplex.EpochConfig{}, err
+		return nil, err
 	}
 
 	wal, err := wal.NewGarbageCollectedWAL(i.Config.WALs, i.Config.WalCreator, &common.WALRetentionReader{}, i.Config.ParameterConfig.WALMaxEntryCount)
 	if err != nil {
-		return simplex.EpochConfig{}, fmt.Errorf("error creating garbage collected wal: %w", err)
+		return nil, fmt.Errorf("error creating garbage collected wal: %w", err)
 	}
 	i.wal = wal
 
 	// We might have crashed right after a sealing block was persisted to storage,
 	// but before the WAL was garbage collected.
 	// In that case, we need to garbage collect the WAL to remove all entries from previous epochs.
 	if err := i.maybeGarbageCollectWAL(lastBlock); err != nil {
-		return simplex.EpochConfig{}, err
+		return nil, err
 	}
 
 	msm, err := metadata.NewStateMachine(&metadata.Config{
@@ -466,15 +466,15 @@ func (i *Instance) createEpochConfig() (simplex.EpochConfig, error) {
 		GetBlock:                        i.cs.RetrieveBlock,
 	})
 	if err != nil {
-		return simplex.EpochConfig{}, fmt.Errorf("error creating metadata state machine: %w", err)
+		return nil, fmt.Errorf("error creating metadata state machine: %w", err)
 	}
 
 	i.msm = msm
 	i.cs.msm = msm
 
 	source, err := simplex.NewRandomSource()
 	if err != nil {
-		return simplex.EpochConfig{}, err
+		return nil, err
 	}
 
 	blockBuilder := &BlockBuilderWaiter{vm: i.Config.VM, msm: msm}
@@ -494,7 +494,7 @@ func (i *Instance) createEpochConfig() (simplex.EpochConfig, error) {
 		},
 	}
 
-	epochConfig := simplex.EpochConfig{
+	ec := simplex.EpochConfig{
 		Epoch:              epochNum,
 		ReplicationEnabled: true,
 		StartTime:          time.Now(),
@@ -516,7 +516,10 @@ func (i *Instance) createEpochConfig() (simplex.EpochConfig, error) {
 		BlockBuilder:               blockBuilder,
 		BlockDeserializer:          &blockDeserializer{vm: i.Config.VM, cs: i.cs},
 	}
-	return epochConfig, nil
+	return &epochConfig{
+		EpochConfig: ec,
+		bbw:         blockBuilder,
+	}, nil
 }
 
 func (i *Instance) maybeGarbageCollectWAL(lastBlock metadata.StateMachineBlock) error {
@@ -650,3 +653,8 @@ func constructValidatorSetFromSealingBlock(lastBlock *ParsedBlock) metadata.Node
 	}
 	return validatorSet
 }
+
+type epochConfig struct {
+	simplex.EpochConfig
+	bbw *BlockBuilderWaiter
+}

instance_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/ava-labs/simplex/avalanchego"
 	"github.com/ava-labs/simplex/common"
 	metadata "github.com/ava-labs/simplex/msm"
+	"github.com/ava-labs/simplex/simplex"
 	"github.com/ava-labs/simplex/testutil"
 	"github.com/ava-labs/simplex/wal"
 
@@ -120,6 +121,71 @@ func TestInstanceMixedNodeType(t *testing.T) {
 	require.Equal(t, secondInstance.Config.ID, latestValidatorID(t, storage))
 }
 
+// emptyVoteRecorder wraps a Broadcaster and signals the first time an empty vote is broadcast.
+type emptyVoteRecorder struct {
+	Broadcaster
+	got chan struct{}
+}
+
+func (r *emptyVoteRecorder) Broadcast(msg *common.Message) {
+	if msg.EmptyVoteMessage != nil {
+		select {
+		case r.got <- struct{}{}:
+		default:
+		}
+	}
+	r.Broadcaster.Broadcast(msg)
+}
+
+func TestEpochInvokesMSMWaitForPendingBlock(t *testing.T) {
+	const basePChainHeight = uint64(1)
+
+	// Two validators, but only one is instantiated. Our node has the smaller ID so it sorts to
+	// index 0 and is a non-leader for round 1 (LeaderForRound picks index 1%2). The other
+	// validator is the round leader but is never created, so no block is ever proposed.
+	var ourID, leaderID [20]byte
+	ourID[0], leaderID[0] = 0x01, 0x02
+	ourNode := common.NodeID(ourID[:])
+
+	require.NotEqual(t, ourNode, simplex.LeaderForRound([]common.NodeID{ourNode, leaderID[:]}, 1)) // ensure the leader is not our node
+
+	validatorSetsAtHeight := map[uint64]metadata.NodeBLSMappings{
+		basePChainHeight: {
+			{NodeID: ourID, BLSKey: []byte{0xaa}, Weight: 1},
+			{NodeID: leaderID, BLSKey: []byte{0xbb}, Weight: 1},
+		},
+	}
+
+	pChain := newTestPlatformChain(basePChainHeight, validatorSetsAtHeight)
+	cops := &testCryptoOps{}
+	genesisBlock := &testInnerBlock{Height_: 0, TS: time.Now(), Payload: []byte("genesis")}
+
+	net := newInMemNetwork(t)
+	t.Cleanup(net.stop)
+
+	storage := newStorageWithGenesis(t, genesisBlock)
+
+	// A paused VM never has a pending block, so its WaitForPendingBlock blocks until its context
+	// is cancelled: the MSM must decide to build on its own for the round to make progress.
+	vm := newTestVM()
+	vm.pause()
+
+	inst := newInstanceWithVM(t, ourNode, storage, net, pChain, cops, genesisBlock, vm)
+
+	// Capture the empty vote the node broadcasts once it gives up waiting for the leader.
+	recorder := &emptyVoteRecorder{Broadcaster: inst.Config.Broadcaster, got: make(chan struct{}, 1)}
+	inst.Config.Broadcaster = recorder
+
+	require.NoError(t, inst.Start(t.Context()))
+	t.Cleanup(inst.Stop)
+
+	select {
+	case <-recorder.got:
+	case <-time.After(10 * time.Second):
+		require.FailNow(t, "node never broadcast an empty vote, so the Epoch did not drive the MSM's WaitForPendingBlock")
+	}
+}
+
 func TestInstanceNonValidatorBootstraps(t *testing.T) {
 	// One node is a validator and progresses the chain by building blocks,
 	// and its weight changes while the chain progresses in 3 different P-chain epoch heights.

msm/msm.go

@@ -271,6 +271,101 @@ func (sm *StateMachine) maybeInitializeApprovalStore(validatorSet NodeBLSMapping
 	return sm.approvalStore
 }
 
+// WaitForPendingBlock waits for either the VM to signal that a block is ready to be built,
+// or for the state machine to determine that a block should be built immediately due to an epoch transition.
+// In the latter case, we only wait up to MaxBlockBuildingWaitTime before returning.
+func (sm *StateMachine) WaitForPendingBlock(ctx context.Context, currentRoundMetadata common.ProtocolMetadata) {
+	if currentRoundMetadata.Seq == 0 {
+		// This shouldn't happen because we never build the genesis block.
+		// Just call the VM's WaitForPendingBlock to avoid blocking the state machine.
+		sm.Logger.Debug("WaitForPendingBlock called with seq 0, which is invalid; forwarding to VM to avoid blocking")
+		sm.BlockBuilder.WaitForPendingBlock(ctx)
+		return
+	}
+
+	// In order to know whether we're transitioning to a new epoch, or should transition to one,
+	// we need to look at the previous block's metadata.
+
+	prevBlockSeq := currentRoundMetadata.Seq - 1
+
+	parentBlock, finalization, err := sm.GetBlock(prevBlockSeq, currentRoundMetadata.Prev)
+	if err != nil {
+		sm.Logger.Debug("WaitForPendingBlock failed to get block", zap.Uint64("seq", prevBlockSeq), zap.Error(err))
+		sm.BlockBuilder.WaitForPendingBlock(ctx)
+		return
+	}
+
+	// In case the parent block is a Telock, we want to look at the sealing block's metadata instead.
+	if parentBlock.Type() == BlockTypeTelock {
+		sealingBlockSeq := parentBlock.Metadata.SimplexEpochInfo.SealingBlockSeq
+		var sealingBlock StateMachineBlock
+		sealingBlock, finalization, err = sm.GetBlock(sealingBlockSeq, common.Digest{})
+		if err != nil {
+			sm.Logger.Debug("Failed retrieving sealing block for previous epoch", zap.Uint64("seq", sealingBlockSeq), zap.Error(err))
+			sm.BlockBuilder.WaitForPendingBlock(ctx)
+			return
+		}
+
+		// If the sealing block isn't finalized, we need to build a Telock immediately to extend the epoch,
+		// so we wait up to MaxBlockBuildingWaitTime and then return.
+		if finalization == nil {
+			ctx, cancel := context.WithTimeout(ctx, sm.MaxBlockBuildingWaitTime)
+			defer cancel()
+			sm.BlockBuilder.WaitForPendingBlock(ctx)
+			return
+		}
+		// Else, err is nil and the sealing block is finalized,
+		// so we can use the sealing block's metadata to determine whether we should build a block immediately or not.
+		parentBlock = sealingBlock
+	}
+
+	currentState := parentBlock.Metadata.SimplexEpochInfo.NextState()
+
+	// We first check if we have obvious signs that we need to build a block immediately:
+	var shouldObviouslyBuildBlockImmediately bool
+
+	switch currentState {
+	case stateFirstSimplexBlock, stateBuildCollectingApprovals:
+		// In case of the first simplex block we need to persist the validator set,
+		// and if we're collecting approvals then we're in the process of an epoch change already.
+		shouldObviouslyBuildBlockImmediately = true
+	case stateBuildBlockEpochSealed:
+		// The parent is a sealing block, whose finalization we already fetched above.
+		// If it isn't finalized we must build a Telock immediately to extend the epoch;
+		// if it is finalized, we fall through to the normal epoch-transition check below.
+		if finalization == nil {
+			shouldObviouslyBuildBlockImmediately = true
+		}
+	default: // Handles stateBuildBlockNormalOp or any unknown state, don't do anything and just exit the switch.
+	}
+
+	// If it's obvious that we should build a block immediately,
+	// we wait up to MaxBlockBuildingWaitTime and then return immediately.
+	if shouldObviouslyBuildBlockImmediately {
+		ctx, cancel := context.WithTimeout(ctx, sm.MaxBlockBuildingWaitTime)
+		defer cancel()
+		sm.BlockBuilder.WaitForPendingBlock(ctx)
+		return
+	}
+
+	// Otherwise, we might need to build a block if we detect that we should transition to a new epoch,
+	// so we initialize a blockBuildingDecider and listen while waiting for the VM.
+	// We return when either the VM signals that a block is ready to be built,
+	// or that the blockBuildingDecider detects that we should transition to a new epoch.
+	pChainReferenceHeight := parentBlock.Metadata.SimplexEpochInfo.PChainReferenceHeight
+	if parentBlock.Type() == BlockTypeSealing {
+		// We've moved to a new epoch, so we need to use the next P-chain reference height of the sealing block,
+		// because the P-chain reference height is of the next epoch is inherited from the P-chain reference height of the sealing block.
+		pChainReferenceHeight = parentBlock.Metadata.SimplexEpochInfo.NextPChainReferenceHeight
+	}
+	blockBuildingDecider := sm.createBlockBuildingDecider(pChainReferenceHeight)
+	_, err = blockBuildingDecider.shouldBuildBlock(ctx)
+	if err != nil {
+		sm.Logger.Debug("Error while deciding whether to build a block", zap.Error(err))
+		return
+	}
+}
+
 // BuildBlock constructs the next block on top of the given parent block, and passes in the provided simplex metadata and blacklist.
 func (sm *StateMachine) BuildBlock(ctx context.Context, metadata common.ProtocolMetadata, blacklist common.Blacklist) (*StateMachineBlock, error) {
 	// The zero sequence number is reserved for the genesis block, which should never be built.