fix(replay): install future epoch metadata at boundaries

pkg/replay/alpenglow_engine.go

@@ -53,6 +53,15 @@ func installAlpenglowValidatorSet(consensusEngine consensusengine.Engine, epoch
 	}
 }
 
+// installEpochTransitionAlpenglowValidatorSets installs the validator sets
+// needed both to execute the newly entered epoch and to verify the future epoch
+// whose leader schedule was prepared at the boundary.
+func installEpochTransitionAlpenglowValidatorSets(consensusEngine consensusengine.Engine, newEpoch, leaderScheduleEpoch uint64) {
+	for _, targetEpoch := range epochTransitionTargetEpochs(newEpoch, leaderScheduleEpoch) {
+		installAlpenglowValidatorSet(consensusEngine, targetEpoch)
+	}
+}
+
 // installCachedAlpenglowValidatorSets installs validator sets for every cached
 // epoch (so certs spanning an epoch boundary verify), plus the current epoch.
 func installCachedAlpenglowValidatorSets(consensusEngine consensusengine.Engine, currentEpoch uint64) {

pkg/replay/block.go

@@ -2560,9 +2560,14 @@ func ReplayBlocks(
 				}
 			}
 
-			// Alpenglow: reinstall the BLS validator set for the new epoch.
+			// Alpenglow: install both the newly entered epoch and the future
+			// leader-schedule epoch prepared at this boundary.
 			if consensusEngine != nil {
-				installAlpenglowValidatorSet(consensusEngine, currentEpoch)
+				installEpochTransitionAlpenglowValidatorSets(
+					consensusEngine,
+					currentEpoch,
+					epochSchedule.LeaderScheduleEpoch(block.Slot),
+				)
 			}
 
 			if len(newlyActivatedFeatures) != 0 {

pkg/replay/epoch.go

@@ -46,6 +46,33 @@ func requireDurableEpochBoundaryParent(boundarySlot, parentSlot, durableSlot uin
 		boundarySlot, parentSlot, durableSlot)
 }
 
+// epochTransitionTargetEpochs returns every epoch whose metadata must be ready
+// when crossing a boundary. The epoch being entered and the epoch whose leader
+// schedule is selected at that slot are usually different, but some epoch
+// schedules make them the same.
+func epochTransitionTargetEpochs(newEpoch, leaderScheduleEpoch uint64) []uint64 {
+	epochs := []uint64{newEpoch}
+	if leaderScheduleEpoch != newEpoch {
+		epochs = append(epochs, leaderScheduleEpoch)
+	}
+	return epochs
+}
+
+func prepareEpochTransitionLeaderSchedules(newEpoch, leaderScheduleEpoch uint64, epochSchedule *sealevel.SysvarEpochSchedule, logsDir string) error {
+	for _, targetEpoch := range epochTransitionTargetEpochs(newEpoch, leaderScheduleEpoch) {
+		var err error
+		if len(global.EpochStakesVoteAccts(targetEpoch)) > 0 {
+			_, err = PrepareLeaderScheduleLocal(targetEpoch, epochSchedule, logsDir)
+		} else {
+			_, err = PrepareLeaderScheduleLocalFromVoteCache(targetEpoch, epochSchedule, logsDir)
+		}
+		if err != nil {
+			return fmt.Errorf("prepare leader schedule for epoch %d at epoch transition: %w", targetEpoch, err)
+		}
+	}
+	return nil
+}
+
 // newReplayCtx creates a new ReplayCtx, preferring values from resumeState if available.
 // This ensures resume uses fresh values instead of potentially stale manifest data.
 func newReplayCtx(mithrilState *state.MithrilState, resumeState *ResumeState) (*ReplayCtx, error) {
@@ -241,13 +268,7 @@ func handleEpochTransition(acctsDb *accountsdb.AccountsDb, partitionedEpochRewar
 	t2 := time.Now()
 
 	if global.ManageLeaderSchedule() {
-		if len(global.EpochStakesVoteAccts(newEpoch)) > 0 {
-			_, err = PrepareLeaderScheduleLocal(newEpoch, epochSchedule, "")
-		} else {
-			_, err = PrepareLeaderScheduleLocalFromVoteCache(newEpoch, epochSchedule, "")
-		}
-
-		if err != nil {
+		if err = prepareEpochTransitionLeaderSchedules(newEpoch, leaderScheduleEpoch, epochSchedule, ""); err != nil {
 			panic(err)
 		}
 

pkg/replay/epoch_transition_metadata_test.go

@@ -0,0 +1,138 @@
+package replay
+
+import (
+	"crypto/ed25519"
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/alpenglow"
+	"github.com/Overclock-Validator/mithril/pkg/epochstakes"
+	"github.com/Overclock-Validator/mithril/pkg/global"
+	"github.com/Overclock-Validator/mithril/pkg/sealevel"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEpochTransitionTargetEpochs(t *testing.T) {
+	const (
+		slotsPerEpoch = uint64(54_000)
+		boundarySlot  = uint64(110 * slotsPerEpoch)
+	)
+
+	tests := []struct {
+		name   string
+		offset uint64
+		want   []uint64
+	}{
+		{
+			name:   "standard one epoch offset",
+			offset: slotsPerEpoch,
+			want:   []uint64{110, 111},
+		},
+		{
+			name:   "same epoch is deduplicated",
+			offset: slotsPerEpoch / 2,
+			want:   []uint64{110},
+		},
+		{
+			name:   "nonstandard offset is not hardcoded",
+			offset: 2 * slotsPerEpoch,
+			want:   []uint64{110, 112},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			epochSchedule := &sealevel.SysvarEpochSchedule{
+				SlotsPerEpoch:            slotsPerEpoch,
+				LeaderScheduleSlotOffset: tt.offset,
+			}
+			newEpoch := epochSchedule.GetEpoch(boundarySlot)
+			leaderScheduleEpoch := epochSchedule.LeaderScheduleEpoch(boundarySlot)
+
+			require.Equal(t, tt.want, epochTransitionTargetEpochs(newEpoch, leaderScheduleEpoch))
+		})
+	}
+}
+
+func TestPrepareEpochTransitionLeaderSchedulesPreparesFutureEpoch(t *testing.T) {
+	const slotsPerEpoch = uint64(54_000)
+	epochSchedule := &sealevel.SysvarEpochSchedule{
+		SlotsPerEpoch:            slotsPerEpoch,
+		LeaderScheduleSlotOffset: slotsPerEpoch,
+	}
+	_, nodePubkey := seedEpochTransitionValidator(t, 110, 111)
+
+	err := prepareEpochTransitionLeaderSchedules(110, 111, epochSchedule, t.TempDir())
+	require.NoError(t, err)
+
+	for _, epoch := range []uint64{110, 111} {
+		leader, ok := global.LeaderForSlot(epochSchedule.FirstSlotInEpoch(epoch))
+		require.Truef(t, ok, "leader schedule for epoch %d was not installed", epoch)
+		require.Equal(t, nodePubkey, leader)
+	}
+}
+
+type epochTransitionValidatorSetRecorder struct {
+	fakeEngine
+	installed []alpenglow.ValidatorSet
+}
+
+func (r *epochTransitionValidatorSetRecorder) SetAlpenglowValidatorSet(set alpenglow.ValidatorSet) error {
+	r.installed = append(r.installed, set)
+	return nil
+}
+
+func TestInstallEpochTransitionAlpenglowValidatorSetsIncludesFuture(t *testing.T) {
+	seedEpochTransitionValidator(t, 110, 111)
+	recorder := new(epochTransitionValidatorSetRecorder)
+
+	installEpochTransitionAlpenglowValidatorSets(recorder, 110, 111)
+
+	require.Len(t, recorder.installed, 2)
+	require.Equal(t, uint64(110), recorder.installed[0].Epoch)
+	require.Equal(t, uint64(111), recorder.installed[1].Epoch)
+}
+
+func TestInstallEpochTransitionAlpenglowValidatorSetsDeduplicatesEqualEpoch(t *testing.T) {
+	seedEpochTransitionValidator(t, 110)
+	recorder := new(epochTransitionValidatorSetRecorder)
+
+	installEpochTransitionAlpenglowValidatorSets(recorder, 110, 110)
+
+	require.Len(t, recorder.installed, 1)
+	require.Equal(t, uint64(110), recorder.installed[0].Epoch)
+}
+
+func seedEpochTransitionValidator(t *testing.T, epochs ...uint64) (solana.PublicKey, solana.PublicKey) {
+	t.Helper()
+
+	var voteAccount, nodePubkey solana.PublicKey
+	voteAccount[0] = 0xa1
+	nodePubkey[0] = 0xb1
+
+	privateKey := ed25519.NewKeyFromSeed(make([]byte, ed25519.SeedSize))
+	signer, err := alpenglow.DeriveBLSSigner(privateKey)
+	require.NoError(t, err)
+	compressed := signer.PublicKeyCompressed()
+
+	for _, epoch := range epochs {
+		global.ClearEpochStakes(epoch)
+		global.SetLeaderScheduleForEpoch(epoch, nil)
+
+		blsPubkey := compressed
+		global.PutEpochStakesEntry(epoch, voteAccount, 100, &epochstakes.VoteAccount{
+			NodePubkey:          nodePubkey,
+			BlsPubkeyCompressed: &blsPubkey,
+		})
+		global.PutEpochTotalStake(epoch, 100)
+	}
+
+	t.Cleanup(func() {
+		for _, epoch := range epochs {
+			global.ClearEpochStakes(epoch)
+			global.SetLeaderScheduleForEpoch(epoch, nil)
+		}
+	})
+
+	return voteAccount, nodePubkey
+}