Resolve cold-state migration boundaries using finalized canonical block root

beacon-chain/blockchain/head_test.go

@@ -889,7 +889,7 @@ func TestUpdateHead_noSavedChanges(t *testing.T) {
 
 	bellatrixState, _ := util.DeterministicGenesisStateBellatrix(t, 2)
 	require.NoError(t, beaconDB.SaveState(ctx, bellatrixState, bellatrixBlkRoot))
-	service.cfg.StateGen.SaveFinalizedState(0, bellatrixBlkRoot, bellatrixState)
+	service.cfg.StateGen.SaveFinalizedState(bellatrixBlkRoot, bellatrixState)
 
 	headRoot := service.headRoot()
 	require.Equal(t, [32]byte{}, headRoot)

beacon-chain/blockchain/service.go

@@ -374,7 +374,7 @@ func (s *Service) saveGenesisData(ctx context.Context, genesisState state.Beacon
 	}
 
 	s.originBlockRoot = genesisBlkRoot
-	s.cfg.StateGen.SaveFinalizedState(0 /*slot*/, genesisBlkRoot, genesisState)
+	s.cfg.StateGen.SaveFinalizedState(genesisBlkRoot, genesisState)
 
 	s.cfg.ForkChoiceStore.Lock()
 	defer s.cfg.ForkChoiceStore.Unlock()

beacon-chain/state/stategen/getter_test.go

@@ -42,7 +42,6 @@ func TestStateByRoot_ColdState(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 
 	service := New(beaconDB, doublylinkedtree.New())
-	service.finalizedInfo.slot = 2
 	service.slotsPerArchivedPoint = 1
 
 	b := util.NewBeaconBlock()
@@ -96,7 +95,6 @@ func TestStateByRootIfCachedNoCopy_ColdState(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 
 	service := New(beaconDB, doublylinkedtree.New())
-	service.finalizedInfo.slot = 2
 	service.slotsPerArchivedPoint = 1
 
 	b := util.NewBeaconBlock()
@@ -264,7 +262,6 @@ func TestLoadStateByRoot(t *testing.T) {
 	persistFinalizedStruct := func(r testChain, slot primitives.Slot) {
 		st := r.state(t, slot)
 		r.srv.finalizedInfo.state = st
-		r.srv.finalizedInfo.slot = st.Slot()
 		r.srv.finalizedInfo.root = r.blockRoot(t, slot)
 	}
 

beacon-chain/state/stategen/migrate.go

@@ -16,9 +16,11 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-// MigrateToCold advances the finalized info in between the cold and hot state sections.
-// It moves the recent finalized states from the hot section to the cold section and
-// only preserves the ones that are on archived point.
+// maxFinalizedRootSearch is for bounding the number of iterations
+// when searching for a finalized canonical root below a given slot.
+const maxFinalizedRootSearch = 1024
+
+// MigrateToCold moves finalized states to cold storage and advances the migration cursor.
 func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 	ctx, span := trace.StartSpan(ctx, "stateGen.MigrateToCold")
 	defer span.End()
@@ -33,9 +35,7 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 		return s.migrateToColdHdiff(ctx, fRoot)
 	}
 
-	s.finalizedInfo.lock.RLock()
-	oldFSlot := s.finalizedInfo.slot
-	s.finalizedInfo.lock.RUnlock()
+	oldFSlot := s.migratedSlot
 
 	fBlock, err := s.beaconDB.Block(ctx, fRoot)
 	if err != nil {
@@ -79,16 +79,10 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 			aRoot = cached.root
 			aState = cached.state
 		} else {
-			_, roots, err := s.beaconDB.HighestRootsBelowSlot(ctx, slot)
+			aRoot, err = s.canonicalRootBelowSlot(ctx, slot)
 			if err != nil {
-				return err
+				return fmt.Errorf("canonical root below slot %d not found: %w", slot, err)
 			}
-			// Given the block has been finalized, the db should not have more than one block in a given slot.
-			// We should error out when this happens.
-			if len(roots) != 1 {
-				return errUnknownBlock
-			}
-			aRoot = roots[0]
 			// There's no need to generate the state if the state already exists in the DB.
 			// We can skip saving the state.
 			if !s.beaconDB.HasState(ctx, aRoot) {
@@ -125,17 +119,18 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 		return err
 	}
 	if ok {
-		s.SaveFinalizedState(fSlot, fRoot, fInfo.state)
+		s.SaveFinalizedState(fRoot, fInfo.state)
 	}
+	// The migration can complete without the finalized state being cached. Keep
+	// finalizedInfo coherent and advance its independent migration cursor.
+	s.migratedSlot = fSlot
 
 	return nil
 }
 
 // migrateToColdHdiff saves the state-diffs for slots that are in the state diff tree after finalization
 func (s *State) migrateToColdHdiff(ctx context.Context, fRoot [32]byte) error {
-	s.finalizedInfo.lock.RLock()
-	oldFSlot := s.finalizedInfo.slot
-	s.finalizedInfo.lock.RUnlock()
+	oldFSlot := s.migratedSlot
 	fSlot, err := s.beaconDB.SlotByBlockRoot(ctx, fRoot)
 	if err != nil {
 		return errors.Wrap(err, "could not get slot by block root")
@@ -171,16 +166,10 @@ func (s *State) migrateToColdHdiff(ctx context.Context, fRoot [32]byte) error {
 		} else {
 			// we check for slot+1 because we don't want to treat this slot as a missed block when it's not in cache.
 			// this specifically happens when the state is evicted from the caches in long non finalization.
-			_, roots, err := s.beaconDB.HighestRootsBelowSlot(ctx, slot+1)
+			aRoot, err = s.canonicalRootBelowSlot(ctx, slot+1)
 			if err != nil {
-				return err
+				return fmt.Errorf("canonical root below slot %d not found: %w", slot, err)
 			}
-			// Given the block has been finalized, the db should not have more than one block in a given slot.
-			// We should error out when this happens.
-			if len(roots) != 1 {
-				return errUnknownBlock
-			}
-			aRoot = roots[0]
 			// Different than the legacy MigrateToCold, we need to always get the state even if
 			// the state exists in DB as part of the hot state db, because we need to process slots
 			// to the state diff tree slots.
@@ -217,8 +206,11 @@ func (s *State) migrateToColdHdiff(ctx context.Context, fRoot [32]byte) error {
 		return err
 	}
 	if ok {
-		s.SaveFinalizedState(fSlot, fRoot, fInfo.state)
+		s.SaveFinalizedState(fRoot, fInfo.state)
 	}
+	// The migration can complete without the finalized state being cached. Keep
+	// finalizedInfo coherent and advance its independent migration cursor.
+	s.migratedSlot = fSlot
 	return nil
 }
 
@@ -237,3 +229,23 @@ func (s *State) migrateHotToCold(aRoot [32]byte) {
 	}
 	s.saveHotStateDB.lock.Unlock()
 }
+
+// canonicalRootBelowSlot returns a finalized canonical block immediately before slot.
+func (s *State) canonicalRootBelowSlot(ctx context.Context, slot primitives.Slot) ([32]byte, error) {
+	for attempts := 0; slot > 0 && attempts < maxFinalizedRootSearch; attempts++ {
+		if err := ctx.Err(); err != nil {
+			return [32]byte{}, err
+		}
+		found, roots, err := s.beaconDB.HighestRootsBelowSlot(ctx, slot)
+		if err != nil {
+			return [32]byte{}, fmt.Errorf("highest roots below slot %d: %w", slot, err)
+		}
+		for _, root := range roots {
+			if s.beaconDB.IsFinalizedBlock(ctx, root) {
+				return root, nil
+			}
+		}
+		slot = found
+	}
+	return [32]byte{}, errUnknownBlock
+}