State migration from hot to cold (archived) (#5076)

* Starting

* Test

* Tests

* comments

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/state/stategen/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
         "errors.go",
         "hot.go",
         "log.go",
+        "migrate.go",
         "replay.go",
         "service.go",
     ],
@@ -37,6 +38,7 @@ go_test(
         "cold_test.go",
         "epoch_boundary_root_test.go",
         "hot_test.go",
+        "migrate_test.go",
         "replay_test.go",
         "service_test.go",
     ],
@@ -53,5 +55,6 @@ go_test(
         "@com_github_gogo_protobuf//proto:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
         "@com_github_prysmaticlabs_go_ssz//:go_default_library",
+        "@com_github_sirupsen_logrus//hooks/test:go_default_library",
     ],
 )

beacon-chain/state/stategen/migrate.go

@@ -0,0 +1,99 @@
+package stategen
+
+import (
+	"context"
+	"encoding/hex"
+
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
+	"github.com/prysmaticlabs/prysm/beacon-chain/db/filters"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/sirupsen/logrus"
+	"go.opencensus.io/trace"
+)
+
+// MigrateToCold advances the split point in between the cold and hot state sections.
+// It moves the recent finalized states from the hot section to the cold section and
+// only preserve the ones that's on archived point.
+func (s *State) MigrateToCold(ctx context.Context, finalizedState *state.BeaconState, finalizedRoot [32]byte) error {
+	ctx, span := trace.StartSpan(ctx, "stateGen.MigrateToCold")
+	defer span.End()
+
+	// Verify migration is sensible. The new finalized point must increase the current split slot, and
+	// on an epoch boundary for hot state summary scheme to work.
+	currentSplitSlot := s.splitInfo.slot
+	if currentSplitSlot > finalizedState.Slot() {
+		return nil
+	}
+	if !helpers.IsEpochStart(finalizedState.Slot()) {
+		return nil
+	}
+
+	// Move the states between split slot to finalized slot from hot section to the cold section.
+	filter := filters.NewFilter().SetStartSlot(currentSplitSlot).SetEndSlot(finalizedState.Slot() - 1)
+	blockRoots, err := s.beaconDB.BlockRoots(ctx, filter)
+	if err != nil {
+		return err
+	}
+
+	for _, r := range blockRoots {
+		stateSummary, err := s.beaconDB.StateSummary(ctx, r)
+		if err != nil {
+			return err
+		}
+		if stateSummary == nil || stateSummary.Slot == 0 {
+			continue
+		}
+
+		if stateSummary.Slot%s.slotsPerArchivedPoint == 0 {
+			archivePointIndex := stateSummary.Slot / s.slotsPerArchivedPoint
+			if s.beaconDB.HasState(ctx, r) {
+				hotState, err := s.beaconDB.State(ctx, r)
+				if err != nil {
+					return err
+				}
+				if err := s.beaconDB.SaveArchivedPointState(ctx, hotState.Copy(), archivePointIndex); err != nil {
+					return err
+				}
+			} else {
+				hotState, err := s.ComputeStateUpToSlot(ctx, stateSummary.Slot)
+				if err != nil {
+					return err
+				}
+				if err := s.beaconDB.SaveArchivedPointState(ctx, hotState.Copy(), archivePointIndex); err != nil {
+					return err
+				}
+			}
+			if err := s.beaconDB.SaveArchivedPointRoot(ctx, r, archivePointIndex); err != nil {
+				return err
+			}
+
+			log.WithFields(logrus.Fields{
+				"slot":         stateSummary.Slot,
+				"archiveIndex": archivePointIndex,
+				"root":         hex.EncodeToString(bytesutil.Trunc(r[:])),
+			}).Info("Saved archived point during state migration")
+		}
+
+		if s.beaconDB.HasState(ctx, r) {
+			if err := s.beaconDB.DeleteState(ctx, r); err != nil {
+				return err
+			}
+			log.WithFields(logrus.Fields{
+				"slot": stateSummary.Slot,
+				"root": hex.EncodeToString(bytesutil.Trunc(r[:])),
+			}).Info("Deleted state during migration")
+		}
+
+		s.deleteEpochBoundaryRoot(stateSummary.Slot)
+	}
+
+	// Update the split slot and root.
+	s.splitInfo = &splitSlotAndRoot{slot: finalizedState.Slot(), root: finalizedRoot}
+	log.WithFields(logrus.Fields{
+		"slot": s.splitInfo.slot,
+		"root": hex.EncodeToString(bytesutil.Trunc(s.splitInfo.root[:])),
+	}).Info("Set hot and cold state split point")
+
+	return nil
+}

beacon-chain/state/stategen/migrate_test.go

@@ -0,0 +1,104 @@
+package stategen
+
+import (
+	"context"
+	"testing"
+
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
+	testDB "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
+	logTest "github.com/sirupsen/logrus/hooks/test"
+)
+
+func TestMigrateToCold_NoBlock(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch)
+	if err := service.MigrateToCold(ctx, beaconState, [32]byte{}); err != nil {
+		t.Fatal(err)
+	}
+
+	testutil.AssertLogsContain(t, hook, "Set hot and cold state split point")
+}
+
+func TestMigrateToCold_HigherSplitSlot(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.splitInfo.slot = 2
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(1)
+	if err := service.MigrateToCold(ctx, beaconState, [32]byte{}); err != nil {
+		t.Fatal(err)
+	}
+
+	testutil.AssertLogsDoNotContain(t, hook, "Set hot and cold state split point")
+}
+
+func TestMigrateToCold_NotEpochStart(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch + 1)
+	if err := service.MigrateToCold(ctx, beaconState, [32]byte{}); err != nil {
+		t.Fatal(err)
+	}
+
+	testutil.AssertLogsDoNotContain(t, hook, "Set hot and cold state split point")
+}
+
+func TestMigrateToCold_MigrationCompletes(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch)
+	b := &ethpb.SignedBeaconBlock{
+		Block: &ethpb.BeaconBlock{Slot: 2},
+	}
+	if err := service.beaconDB.SaveBlock(ctx, b); err != nil {
+		t.Fatal(err)
+	}
+	bRoot, _ := ssz.HashTreeRoot(b.Block)
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{Root: bRoot[:], Slot: 2}); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveState(ctx, beaconState, bRoot); err != nil {
+		t.Fatal(err)
+	}
+	service.slotsPerArchivedPoint = 2 // Ensure we can land on archived point.
+
+	if err := service.MigrateToCold(ctx, beaconState, [32]byte{}); err != nil {
+		t.Fatal(err)
+	}
+
+	if !service.beaconDB.HasArchivedPoint(ctx, 1) {
+		t.Error("Did not preserve archived point")
+	}
+
+	testutil.AssertLogsContain(t, hook, "Saved archived point during state migration")
+	testutil.AssertLogsContain(t, hook, "Deleted state during migration")
+	testutil.AssertLogsContain(t, hook, "Set hot and cold state split point")
+}

beacon-chain/state/stategen/service.go

@@ -17,6 +17,14 @@ type State struct {
 	epochBoundarySlotToRoot map[uint64][32]byte
 	epochBoundaryLock       sync.RWMutex
 	hotStateCache           *cache.HotStateCache
+	splitInfo               *splitSlotAndRoot
+}
+
+// This tracks the split point. The point where slot and the block root of
+// cold and hot sections of the DB splits.
+type splitSlotAndRoot struct {
+	slot uint64
+	root [32]byte
 }
 
 // New returns a new state management object.
@@ -25,6 +33,7 @@ func New(db db.NoHeadAccessDatabase) *State {
 		beaconDB:                db,
 		epochBoundarySlotToRoot: make(map[uint64][32]byte),
 		hotStateCache:           cache.NewHotStateCache(),
+		splitInfo:               &splitSlotAndRoot{slot: 0, root: params.BeaconConfig().ZeroHash},
 	}
 }