From 537fe3f72125a748f5c5266cf76ba324c85dfdeb Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Mon, 10 Aug 2020 16:35:28 -0700
Subject: [PATCH] Fix loading historical blocks to replay cold state (#6956)

* Replay: cold state only uses finalized blocks
* Replay: regression test
* Merge refs/heads/master into fix-finalized-replay
---
 beacon-chain/state/stategen/replay.go      | 32 ++++++++++++++++++++--
 beacon-chain/state/stategen/replay_test.go | 23 ++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/beacon-chain/state/stategen/replay.go b/beacon-chain/state/stategen/replay.go
index ec26abf68..291a950e7 100644
--- a/beacon-chain/state/stategen/replay.go
+++ b/beacon-chain/state/stategen/replay.go
@@ -116,6 +116,7 @@ func (s *State) LoadBlocks(ctx context.Context, startSlot uint64, endSlot uint64
 	if err != nil {
 		return nil, err
 	}
+
 	blockRoots, err := s.beaconDB.BlockRoots(ctx, filter)
 	if err != nil {
 		return nil, err
@@ -327,20 +328,21 @@ func (s *State) processStateUpTo(ctx context.Context, state *state.BeaconState,
 	if state.Slot() >= slot {
 		return state, nil
 	}
-
-	lastBlockRoot, lastBlockSlot, err := s.lastSavedBlock(ctx, slot)
+	_, lastBlockSlot, err := s.lastSavedBlock(ctx, slot)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get last saved block")
 	}
+
 	// Short circuit if no block was saved, replay using slots only.
 	if lastBlockSlot == 0 {
 		return s.ReplayBlocks(ctx, state, []*ethpb.SignedBeaconBlock{}, slot)
 	}
 
-	blks, err := s.LoadBlocks(ctx, state.Slot()+1, lastBlockSlot, lastBlockRoot)
+	blks, err := s.loadFinalizedBlocks(ctx, state.Slot()+1, lastBlockSlot)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not load blocks")
 	}
+
 	state, err = s.ReplayBlocks(ctx, state, blks, slot)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not replay blocks")
@@ -348,3 +350,27 @@ func (s *State) processStateUpTo(ctx context.Context, state *state.BeaconState,
 
 	return state, nil
 }
+
+// Given the start slot and the end slot, this returns the finalized beacon blocks in between.
+// Since hot states don't have finalized blocks, this should ONLY be used for replaying cold state.
+func (s *State) loadFinalizedBlocks(ctx context.Context, startSlot uint64, endSlot uint64) ([]*ethpb.SignedBeaconBlock, error) {
+	f := filters.NewFilter().SetStartSlot(startSlot).SetEndSlot(endSlot)
+	bs, err := s.beaconDB.Blocks(ctx, f)
+	if err != nil {
+		return nil, err
+	}
+	bRoots, err := s.beaconDB.BlockRoots(ctx, f)
+	if err != nil {
+		return nil, err
+	}
+	if len(bs) != len(bRoots) {
+		return nil, errors.New("length of blocks and roots don't match")
+	}
+	fbs := make([]*ethpb.SignedBeaconBlock, 0, len(bs))
+	for i := len(bs) - 1; i >= 0; i-- {
+		if s.beaconDB.IsFinalizedBlock(ctx, bRoots[i]) {
+			fbs = append(fbs, bs[i])
+		}
+	}
+	return fbs, nil
+}
diff --git a/beacon-chain/state/stategen/replay_test.go b/beacon-chain/state/stategen/replay_test.go
index 3da1bd2eb..ff3178658 100644
--- a/beacon-chain/state/stategen/replay_test.go
+++ b/beacon-chain/state/stategen/replay_test.go
@@ -697,3 +697,26 @@ func tree4(db db.Database, genesisRoot []byte) ([][32]byte, []*ethpb.SignedBeaco
 
 	return [][32]byte{r0, r21, r22, r23, r24}, returnedBlocks, nil
 }
+
+func TestLoadFinalizedBlocks(t *testing.T) {
+	db, _ := testDB.SetupDB(t)
+	ctx := context.Background()
+	s := &State{
+		beaconDB: db,
+	}
+	gBlock := &ethpb.SignedBeaconBlock{}
+	gRoot, err := stateutil.BlockRoot(gBlock.Block)
+	require.NoError(t, err)
+	require.NoError(t, db.SaveBlock(ctx, gBlock))
+	roots, _, err := tree1(db, gRoot[:])
+	require.NoError(t, err)
+
+	filteredBlocks, err := s.loadFinalizedBlocks(ctx, 0, 8)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(filteredBlocks))
+	require.NoError(t, db.SaveStateSummary(ctx, &pb.StateSummary{Root: roots[8][:]}))
+
+	require.NoError(t, s.beaconDB.SaveFinalizedCheckpoint(ctx, &ethpb.Checkpoint{Root: roots[8][:]}))
+	filteredBlocks, err = s.loadFinalizedBlocks(ctx, 0, 8)
+	require.Equal(t, 10, len(filteredBlocks))
+}