Migrate to cold with state diffs

beacon-chain/db/iface/interface.go

@@ -89,13 +89,15 @@ 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
 	DeleteState(ctx context.Context, blockRoot [32]byte) error
 	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,9 +5,12 @@ 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/encoding/bytesutil"
 	"github.com/OffchainLabs/prysm/v7/monitoring/tracing/trace"
+	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
 
@@ -24,6 +27,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()
@@ -82,19 +89,7 @@ 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
 			}
 
@@ -120,3 +115,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. 

tools/analyzers/recursivelock/analyzer.go

@@ -525,6 +525,9 @@ func hasNestedlock(fullRLockSelector *selIdentList, goPos token.Pos, compareMap
 		if node == (*ast.FuncDecl)(nil) {
 			return ""
 		} else if castedNode, ok := node.(*ast.FuncDecl); ok && castedNode.Recv != nil {
+			if len(castedNode.Recv.List) == 0 || len(castedNode.Recv.List[0].Names) == 0 {
+				return ""
+			}
 			recv = castedNode.Recv.List[0].Names[0]
 			rLockSelector.changeRoot(recv, pass.TypesInfo.ObjectOf(recv))
 		}