Migrate to cold with state diffs

This PR adds the logic to migrate to cold when the database has the
hdiff feature. The main difference is that the boundary states have to
have the right slot therefore they need to be advanced and aren't
necessarily the post-state of a given beacon block root.

beacon-chain/db/iface/interface.go

@@ -89,6 +89,7 @@ type NoHeadAccessDatabase interface {
 	SaveBlocks(ctx context.Context, blocks []interfaces.ReadOnlySignedBeaconBlock) error
 	SaveROBlocks(ctx context.Context, blks []blocks.ROBlock, cache bool) error
 	SaveGenesisBlockRoot(ctx context.Context, blockRoot [32]byte) error
+	SlotByBlockRoot(context.Context, [32]byte) (primitives.Slot, error)
 	// State related methods.
 	SaveState(ctx context.Context, state state.ReadOnlyBeaconState, blockRoot [32]byte) error
 	SaveStates(ctx context.Context, states []state.ReadOnlyBeaconState, blockRoots [][32]byte) error
@@ -96,6 +97,7 @@ type NoHeadAccessDatabase interface {
 	DeleteStates(ctx context.Context, blockRoots [][32]byte) error
 	SaveStateSummary(ctx context.Context, summary *ethpb.StateSummary) error
 	SaveStateSummaries(ctx context.Context, summaries []*ethpb.StateSummary) error
+	SlotInDiffTree(primitives.Slot) (uint64, int, error)
 	// Checkpoint operations.
 	SaveJustifiedCheckpoint(ctx context.Context, checkpoint *ethpb.Checkpoint) error
 	SaveFinalizedCheckpoint(ctx context.Context, checkpoint *ethpb.Checkpoint) error

beacon-chain/db/kv/state_diff.go

@@ -23,6 +23,16 @@ const (
 	The data at level 0 is saved every 2**exponent[0] slots and always contains a full state snapshot that is used as a base for the delta saved at other levels.
 */
 
+// SlotInDiffTree returns whether the given slot is a saving point in the diff tree.
+// It it is, it also returns the offset and level in the tree.
+func (s *Store) SlotInDiffTree(slot primitives.Slot) (uint64, int, error) {
+	offset := s.getOffset()
+	if uint64(slot) < offset {
+		return 0, -1, ErrSlotBeforeOffset
+	}
+	return offset, computeLevel(offset, slot), nil
+}
+
 // saveStateByDiff takes a state and decides between saving a full state snapshot or a diff.
 func (s *Store) saveStateByDiff(ctx context.Context, st state.ReadOnlyBeaconState) error {
 	_, span := trace.StartSpan(ctx, "BeaconDB.saveStateByDiff")
@@ -33,13 +43,10 @@ func (s *Store) saveStateByDiff(ctx context.Context, st state.ReadOnlyBeaconStat
 	}
 
 	slot := st.Slot()
-	offset := s.getOffset()
-	if uint64(slot) < offset {
-		return ErrSlotBeforeOffset
+	offset, lvl, err := s.SlotInDiffTree(slot)
+	if err != nil {
+		return errors.Wrap(err, "could not determine if slot is in diff tree")
 	}
-
-	// Find the level to save the state.
-	lvl := computeLevel(offset, slot)
 	if lvl == -1 {
 		return nil
 	}

beacon-chain/state/stategen/BUILD.bazel

@@ -29,6 +29,7 @@ go_library(
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/sync/backfill/coverage:go_default_library",
         "//cache/lru:go_default_library",
+        "//config/features:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",
         "//consensus-types/interfaces:go_default_library",

beacon-chain/state/stategen/migrate.go

@@ -5,10 +5,13 @@ import (
 	"encoding/hex"
 	"fmt"
 
+	"github.com/OffchainLabs/prysm/v7/beacon-chain/core/transition"
 	"github.com/OffchainLabs/prysm/v7/beacon-chain/state"
+	"github.com/OffchainLabs/prysm/v7/config/features"
 	"github.com/OffchainLabs/prysm/v7/consensus-types/primitives"
 	"github.com/OffchainLabs/prysm/v7/encoding/bytesutil"
 	"github.com/OffchainLabs/prysm/v7/monitoring/tracing/trace"
+	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
 
@@ -25,6 +28,10 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 	s.migrationLock.Lock()
 	defer s.migrationLock.Unlock()
 
+	if features.Get().EnableStateDiff {
+		return s.migrateToColdHdiff(ctx, fRoot)
+	}
+
 	s.finalizedInfo.lock.RLock()
 	oldFSlot := s.finalizedInfo.slot
 	s.finalizedInfo.lock.RUnlock()
@@ -90,21 +97,8 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 				}
 			}
 		}
-
 		if s.beaconDB.HasState(ctx, aRoot) {
-			// If you are migrating a state and its already part of the hot state cache saved to the db,
-			// you can just remove it from the hot state cache as it becomes redundant.
-			s.saveHotStateDB.lock.Lock()
-			roots := s.saveHotStateDB.blockRootsOfSavedStates
-			for i := range roots {
-				if aRoot == roots[i] {
-					s.saveHotStateDB.blockRootsOfSavedStates = append(roots[:i], roots[i+1:]...)
-					// There shouldn't be duplicated roots in `blockRootsOfSavedStates`.
-					// Break here is ok.
-					break
-				}
-			}
-			s.saveHotStateDB.lock.Unlock()
+			s.migrateHotToCold(aRoot)
 			continue
 		}
 
@@ -129,3 +123,103 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 
 	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()
+	fSlot, err := s.beaconDB.SlotByBlockRoot(ctx, fRoot)
+	if err != nil {
+		return errors.Wrap(err, "could not get slot by block root")
+	}
+	for slot := oldFSlot; slot < fSlot; slot++ {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+		_, lvl, err := s.beaconDB.SlotInDiffTree(slot)
+		if err != nil {
+			log.WithError(err).Errorf("could not determine if slot %d is in diff tree", slot)
+			continue
+		}
+		if lvl == -1 {
+			continue
+		}
+		// The state needs to be saved.
+		// Try the epoch boundary cache first.
+		cached, exists, err := s.epochBoundaryStateCache.getBySlot(slot)
+		if err != nil {
+			log.WithError(err).Errorf("could not get epoch boundary state for slot %d", slot)
+			cached = nil
+			exists = false
+		}
+		var aRoot [32]byte
+		var aState state.BeaconState
+		if exists {
+			aRoot = cached.root
+			aState = cached.state
+		} else {
+			_, roots, err := s.beaconDB.HighestRootsBelowSlot(ctx, slot)
+			if err != nil {
+				return 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.
+			aState, err = s.StateByRoot(ctx, aRoot)
+			if err != nil {
+				return err
+			}
+		}
+		if s.beaconDB.HasState(ctx, aRoot) {
+			s.migrateHotToCold(aRoot)
+			continue
+		}
+		// advance slots to the target slot
+		if aState.Slot() < slot {
+			aState, err = transition.ProcessSlots(ctx, aState, slot)
+			if err != nil {
+				return errors.Wrapf(err, "could not process slots to slot %d", slot)
+			}
+		}
+		if err := s.beaconDB.SaveState(ctx, aState, aRoot); err != nil {
+			return err
+		}
+		log.WithFields(
+			logrus.Fields{
+				"slot": aState.Slot(),
+				"root": fmt.Sprintf("%#x", aRoot),
+			}).Info("Saved state in DB")
+	}
+	// Update finalized info in memory.
+	fInfo, ok, err := s.epochBoundaryStateCache.getByBlockRoot(fRoot)
+	if err != nil {
+		return err
+	}
+	if ok {
+		s.SaveFinalizedState(fSlot, fRoot, fInfo.state)
+	}
+	return nil
+}
+
+func (s *State) migrateHotToCold(aRoot [32]byte) {
+	// If you are migrating a state and its already part of the hot state cache saved to the db,
+	// you can just remove it from the hot state cache as it becomes redundant.
+	s.saveHotStateDB.lock.Lock()
+	roots := s.saveHotStateDB.blockRootsOfSavedStates
+	for i := range roots {
+		if aRoot == roots[i] {
+			s.saveHotStateDB.blockRootsOfSavedStates = append(roots[:i], roots[i+1:]...)
+			// There shouldn't be duplicated roots in `blockRootsOfSavedStates`.
+			// Break here is ok.
+			break
+		}
+	}
+	s.saveHotStateDB.lock.Unlock()
+}

changelog/potuz_hdiff_migrate_to_cold.md

@@ -0,0 +1,3 @@
+### Added
+
+- Migrate to cold with the hdiff feature.