Initialize chain info w/o clear db (#3365)

* Initialize chain info upon restart

* Test

beacon-chain/blockchain/forkchoice/service.go

@@ -35,7 +35,7 @@ type Store struct {
 	justifiedCheckpt *ethpb.Checkpoint
 	finalizedCheckpt *ethpb.Checkpoint
 	checkpointState  *cache.CheckpointStateCache
-	lock         sync.RWMutex
+	lock             sync.RWMutex
 }
 
 // NewForkChoiceService instantiates a new service instance that will

beacon-chain/blockchain/service.go

@@ -99,6 +99,9 @@ func (c *ChainService) Start() {
 	if beaconState != nil {
 		log.Info("Beacon chain data already exists, starting service")
 		c.genesisTime = time.Unix(int64(beaconState.GenesisTime), 0)
+		if err := c.initializeChainInfo(c.ctx); err != nil {
+			log.Fatalf("Could not set up chain info: %v", err)
+		}
 		justifiedCheckpoint, err := c.beaconDB.JustifiedCheckpoint(c.ctx)
 		if err != nil {
 			log.Fatalf("Could not get justified checkpoint: %v", err)
@@ -107,7 +110,6 @@ func (c *ChainService) Start() {
 		if err != nil {
 			log.Fatalf("Could not get finalized checkpoint: %v", err)
 		}
-
 		if err := c.forkChoiceStore.GenesisStore(c.ctx, justifiedCheckpoint, finalizedCheckpoint); err != nil {
 			log.Fatalf("Could not start fork choice service: %v", err)
 		}
@@ -257,3 +259,26 @@ func (c *ChainService) saveGenesisValidators(ctx context.Context, s *pb.BeaconSt
 	}
 	return nil
 }
+
+// This gets called to initialize chain info variables using the head stored in DB
+func (c *ChainService) initializeChainInfo(ctx context.Context) error {
+	headBlock, err := c.beaconDB.HeadBlock(ctx)
+	if err != nil {
+		return errors.Wrap(err, "could not get head block in db")
+	}
+	headState, err := c.beaconDB.HeadState(ctx)
+	if err != nil {
+		return errors.Wrap(err, "could not get head state in db")
+	}
+	c.headSlot = headBlock.Slot
+	c.headBlock = headBlock
+	c.headState = headState
+
+	headRoot, err := ssz.SigningRoot(headBlock)
+	if err != nil {
+		return errors.Wrap(err, "could not sign root on head block")
+	}
+	c.canonicalRoots[c.headSlot] = headRoot[:]
+
+	return nil
+}

beacon-chain/blockchain/service_test.go

@@ -1,11 +1,13 @@
 package blockchain
 
 import (
+	"bytes"
 	"context"
 	"encoding/hex"
 	"errors"
 	"io/ioutil"
 	"math/big"
+	"reflect"
 	"testing"
 	"time"
 
@@ -340,3 +342,38 @@ func TestChainService_InitializeBeaconChain(t *testing.T) {
 		t.Error("Canonical root for slot 0 can't be nil after initialize beacon chain")
 	}
 }
+
+func TestChainService_InitializeChainInfo(t *testing.T) {
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	ctx := context.Background()
+
+	headBlock := &ethpb.BeaconBlock{Slot: 1}
+	headState := &pb.BeaconState{Slot: 1}
+	headRoot, _ := ssz.SigningRoot(headBlock)
+	if err := db.SaveState(ctx, headState, headRoot); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SaveBlock(ctx, headBlock); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SaveHeadBlockRoot(ctx, headRoot); err != nil {
+		t.Fatal(err)
+	}
+	c := &ChainService{beaconDB: db, canonicalRoots: make(map[uint64][]byte)}
+	if err := c.initializeChainInfo(ctx); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(c.HeadBlock(), headBlock) {
+		t.Error("head block incorrect")
+	}
+	if !reflect.DeepEqual(c.HeadState(), headState) {
+		t.Error("head block incorrect")
+	}
+	if headBlock.Slot != c.HeadSlot() {
+		t.Error("head slot incorrect")
+	}
+	if !bytes.Equal(headRoot[:], c.HeadRoot()) {
+		t.Error("head slot incorrect")
+	}
+}