Save cold state (#5077)

beacon-chain/state/stategen/cold.go

@@ -2,12 +2,40 @@ package stategen
 
 import (
 	"context"
+	"encoding/hex"
 
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/sirupsen/logrus"
 	"go.opencensus.io/trace"
 )
 
+// This saves a pre finalized beacon state in the cold section of the DB. The returns an error
+// and not store anything if the state does not lie on an archive point boundary.
+func (s *State) saveColdState(ctx context.Context, blockRoot [32]byte, state *state.BeaconState) error {
+	ctx, span := trace.StartSpan(ctx, "stateGen.saveColdState")
+	defer span.End()
+
+	if state.Slot()%s.slotsPerArchivedPoint != 0 {
+		return errSlotNonArchivedPoint
+	}
+
+	archivedPointIndex := state.Slot() / s.slotsPerArchivedPoint
+	if err := s.beaconDB.SaveArchivedPointState(ctx, state, archivedPointIndex); err != nil {
+		return err
+	}
+	if err := s.beaconDB.SaveArchivedPointRoot(ctx, blockRoot, archivedPointIndex); err != nil {
+		return err
+	}
+
+	log.WithFields(logrus.Fields{
+		"slot":      state.Slot(),
+		"blockRoot": hex.EncodeToString(bytesutil.Trunc(blockRoot[:]))}).Info("Saved full state on archived point")
+
+	return nil
+}
+
 // Given the archive index, this returns the archived cold state in the DB.
 // If the archived state does not exist in the state, it'll compute it and save it.
 func (s *State) archivedPointByIndex(ctx context.Context, archiveIndex uint64) (*state.BeaconState, error) {

beacon-chain/state/stategen/cold_test.go

@@ -11,6 +11,53 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 )
 
+func TestSaveColdState_NonArchivedPoint(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.slotsPerArchivedPoint = 2
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(1)
+
+	if err := service.saveColdState(ctx, [32]byte{}, beaconState); err != errSlotNonArchivedPoint {
+		t.Error("Did not get wanted error")
+	}
+}
+
+func TestSaveColdState_CanSave(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.slotsPerArchivedPoint = 1
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(1)
+
+	r := [32]byte{'a'}
+	if err := service.saveColdState(ctx, r, beaconState); err != nil {
+		t.Fatal(err)
+	}
+
+	if !service.beaconDB.HasArchivedPoint(ctx, 1) {
+		t.Error("Did not save cold state")
+	}
+
+	if service.beaconDB.ArchivedPointRoot(ctx, 1) != r {
+		t.Error("Did not get wanted root")
+	}
+
+	receivedState, err := service.beaconDB.ArchivedPointState(ctx, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.Equal(receivedState.InnerStateUnsafe(), beaconState.InnerStateUnsafe()) {
+		t.Error("Did not get wanted state")
+	}
+}
+
 func TestArchivedPointByIndex_HasPoint(t *testing.T) {
 	ctx := context.Background()
 	db := testDB.SetupDB(t)

beacon-chain/sync/initial-sync/blocks_fetcher.go

@@ -74,7 +74,7 @@ func newBlocksFetcher(ctx context.Context, cfg *blocksFetcherConfig) *blocksFetc
 	rateLimiter := leakybucket.NewCollector(
 		allowedBlocksPerSecond, /* rate */
 		allowedBlocksPerSecond, /* capacity */
-		false                   /* deleteEmptyBuckets */)
+		false /* deleteEmptyBuckets */)
 
 	return &blocksFetcher{
 		ctx:            ctx,