Methods to retrieves last saved state and block for stategen pkg (#5005)

* Added last saved block and state * Genesis tests * Gaz * Added state tests Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
prysmaticlabs · Mar 5, 2020 · aebc883 · aebc883
1 parent f3dc113
commit aebc883
Show file tree

Hide file tree

Showing 4 changed files with 265 additions and 2 deletions.
diff --git a/beacon-chain/state/stategen/BUILD.bazel b/beacon-chain/state/stategen/BUILD.bazel
@@ -20,6 +20,7 @@ go_library(
         "//shared/featureconfig:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
+        "@com_github_prysmaticlabs_go_ssz//:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",
         "@io_opencensus_go//trace:go_default_library",
     ],

diff --git a/beacon-chain/state/stategen/replay.go b/beacon-chain/state/stategen/replay.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
 	transition "github.com/prysmaticlabs/prysm/beacon-chain/core/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/filters"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
@@ -110,7 +111,7 @@ func executeStateTransitionStateGen(
 		return nil, ctx.Err()
 	}
 	if signed == nil || signed.Block == nil {
-		return nil, errors.New("nil block")
+		return nil, errUnknownBlock
 	}
 
 	ctx, span := trace.StartSpan(ctx, "beacon-chain.ChainService.ExecuteStateTransitionStateGen")
@@ -142,7 +143,7 @@ func processSlotsStateGen(ctx context.Context, state *stateTrie.BeaconState, slo
 	ctx, span := trace.StartSpan(ctx, "beacon-chain.ChainService.ProcessSlotsStateGen")
 	defer span.End()
 	if state == nil {
-		return nil, errors.New("nil state")
+		return nil, errUnknownState
 	}
 
 	if state.Slot() > slot {
@@ -170,3 +171,86 @@ func processSlotsStateGen(ctx context.Context, state *stateTrie.BeaconState, slo
 
 	return state, nil
 }
+
+// This finds the last saved block in DB from searching backwards from input slot,
+// it returns the block root and the slot of the block.
+// This is used by both hot and cold state management.
+func (s *State) lastSavedBlock(ctx context.Context, slot uint64) ([32]byte, uint64, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.lastSavedBlock")
+	defer span.End()
+
+	// Handle the genesis case where the input slot is 0.
+	if slot == 0 {
+		gRoot, err := s.genesisRoot(ctx)
+		if err != nil {
+			return [32]byte{}, 0, err
+		}
+		return gRoot, 0, nil
+	}
+
+	// Lower bound set as last archived slot is a reasonable assumption given
+	// block is saved at an archived point.
+	filter := filters.NewFilter().SetStartSlot(s.lastArchivedSlot).SetEndSlot(slot)
+	rs, err := s.beaconDB.BlockRoots(ctx, filter)
+	if err != nil {
+		return [32]byte{}, 0, err
+	}
+	if len(rs) == 0 {
+		return [32]byte{}, 0, errors.New("block root has 0 length")
+	}
+	lastRoot := rs[len(rs)-1]
+
+	b, err := s.beaconDB.Block(ctx, lastRoot)
+	if err != nil {
+		return [32]byte{}, 0, err
+	}
+	if b == nil || b.Block == nil {
+		return [32]byte{}, 0, errUnknownBlock
+	}
+
+	return lastRoot, b.Block.Slot, nil
+}
+
+// This finds the last saved state in DB from searching backwards from input slot,
+// it returns the block root of the block which was used to produce the state.
+// This is used by both hot and cold state management.
+func (s *State) lastSavedState(ctx context.Context, slot uint64) ([32]byte, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.lastSavedState")
+	defer span.End()
+
+	// Handle the genesis case where the input slot is 0.
+	if slot == 0 {
+		gRoot, err := s.genesisRoot(ctx)
+		if err != nil {
+			return [32]byte{}, err
+		}
+		return gRoot, nil
+	}
+
+	// Lower bound set as last archived slot is a reasonable assumption given
+	// state is saved at an archived point.
+	filter := filters.NewFilter().SetStartSlot(s.lastArchivedSlot).SetEndSlot(slot)
+	rs, err := s.beaconDB.BlockRoots(ctx, filter)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	if len(rs) == 0 {
+		return [32]byte{}, errors.New("block root has 0 length")
+	}
+	for i := len(rs) - 1; i >= 0; i-- {
+		// Stop until a state is saved.
+		if s.beaconDB.HasState(ctx, rs[i]) {
+			return rs[i], nil
+		}
+	}
+	return [32]byte{}, errUnknownState
+}
+
+// This returns the genesis root.
+func (s *State) genesisRoot(ctx context.Context) ([32]byte, error) {
+	b, err := s.beaconDB.GenesisBlock(ctx)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	return ssz.HashTreeRoot(b.Block)
+}
diff --git a/beacon-chain/state/stategen/replay_test.go b/beacon-chain/state/stategen/replay_test.go
@@ -282,6 +282,183 @@ func TestLoadBlocks_BadStart(t *testing.T) {
 	}
 }
 
+func TestLastSavedBlock_Genesis(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	gBlk := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	gRoot, err := ssz.HashTreeRoot(gBlk.Block)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.beaconDB.SaveBlock(ctx, gBlk); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.beaconDB.SaveGenesisBlockRoot(ctx, gRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	savedRoot, savedSlot, err := s.lastSavedBlock(ctx, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if savedSlot != 0 {
+		t.Error("Did not save genesis slot")
+	}
+	if savedRoot != savedRoot {
+		t.Error("Did not save genesis root")
+	}
+}
+
+func TestLastSavedBlock_CanGet(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	b1 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 5}}
+	if err := s.beaconDB.SaveBlock(ctx, b1); err != nil {
+		t.Fatal(err)
+	}
+	b2 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 10}}
+	if err := s.beaconDB.SaveBlock(ctx, b2); err != nil {
+		t.Fatal(err)
+	}
+	b3 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 20}}
+	if err := s.beaconDB.SaveBlock(ctx, b3); err != nil {
+		t.Fatal(err)
+	}
+
+	savedRoot, savedSlot, err := s.lastSavedBlock(ctx, s.lastArchivedSlot+100)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if savedSlot != s.lastArchivedSlot+20 {
+		t.Error("Did not save correct slot")
+	}
+	wantedRoot, _ := ssz.HashTreeRoot(b3.Block)
+	if savedRoot != wantedRoot {
+		t.Error("Did not save correct root")
+	}
+}
+
+func TestLastSavedBlock_OutOfRange(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	b1 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 127}}
+	if err := s.beaconDB.SaveBlock(ctx, b1); err != nil {
+		t.Fatal(err)
+	}
+
+	_, _, err := s.lastSavedBlock(ctx, s.lastArchivedSlot+1)
+	if err.Error() != "block root has 0 length" {
+		t.Error("Did not get wanted error")
+	}
+}
+
+func TestLastSavedState_Genesis(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	gBlk := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	gRoot, err := ssz.HashTreeRoot(gBlk.Block)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.beaconDB.SaveBlock(ctx, gBlk); err != nil {
+		t.Fatal(err)
+	}
+	if err := s.beaconDB.SaveGenesisBlockRoot(ctx, gRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	savedRoot, err := s.lastSavedState(ctx, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if savedRoot != savedRoot {
+		t.Error("Did not save genesis root")
+	}
+}
+
+func TestLastSavedState_CanGet(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	b1 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 5}}
+	if err := s.beaconDB.SaveBlock(ctx, b1); err != nil {
+		t.Fatal(err)
+	}
+	b2 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 10}}
+	if err := s.beaconDB.SaveBlock(ctx, b2); err != nil {
+		t.Fatal(err)
+	}
+	b2Root, _ := ssz.HashTreeRoot(b2.Block)
+	st, err := stateTrie.InitializeFromProtoUnsafe(&pb.BeaconState{Slot: s.lastArchivedSlot + 10})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.beaconDB.SaveState(ctx, st, b2Root); err != nil {
+		t.Fatal(err)
+	}
+	b3 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: s.lastArchivedSlot + 20}}
+	if err := s.beaconDB.SaveBlock(ctx, b3); err != nil {
+		t.Fatal(err)
+	}
+
+	savedRoot, err := s.lastSavedState(ctx, s.lastArchivedSlot+100)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if savedRoot != b2Root {
+		t.Error("Did not save correct root")
+	}
+}
+
+func TestLastSavedState_OutOfRange(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+	s := &State{
+		beaconDB:         db,
+		lastArchivedSlot: 128,
+	}
+
+	b1 := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 127}}
+	if err := s.beaconDB.SaveBlock(ctx, b1); err != nil {
+		t.Fatal(err)
+	}
+
+	_, err := s.lastSavedState(ctx, s.lastArchivedSlot+1)
+	if err.Error() != "block root has 0 length" {
+		t.Error("Did not get wanted error")
+	}
+}
+
 // tree1 constructs the following tree:
 // B0 - B1 - - B3 -- B5
 //        \- B2 -- B4 -- B6 ----- B8

diff --git a/beacon-chain/state/stategen/service.go b/beacon-chain/state/stategen/service.go
@@ -10,6 +10,7 @@ import (
 // logic of maintaining both hot and cold states in DB.
 type State struct {
 	beaconDB                db.NoHeadAccessDatabase
+	lastArchivedSlot        uint64
 	epochBoundarySlotToRoot map[uint64][32]byte
 	epochBoundaryLock       sync.RWMutex
 }