From 16a0c9f463b1da6a1353d1bdef2a41179cb36fe7 Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Sat, 8 Feb 2020 21:08:21 -0800
Subject: [PATCH] Don't save nil head state (#4799)

* Don't save nil head state
* Update head
---
 beacon-chain/blockchain/chain_info.go | 7 ++++++-
 beacon-chain/blockchain/head.go       | 3 +++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/beacon-chain/blockchain/chain_info.go b/beacon-chain/blockchain/chain_info.go
index af09118b3..ba93a79ea 100644
--- a/beacon-chain/blockchain/chain_info.go
+++ b/beacon-chain/blockchain/chain_info.go
@@ -155,7 +155,12 @@ func (s *Service) HeadState(ctx context.Context) (*state.BeaconState, error) {
 	defer s.headLock.RUnlock()
 
 	if s.headState == nil {
-		return s.beaconDB.HeadState(ctx)
+		headState, err := s.beaconDB.HeadState(ctx)
+		if err != nil {
+			return nil, err
+		}
+		s.headState = headState
+		return headState, nil
 	}
 
 	return s.headState.Copy(), nil
diff --git a/beacon-chain/blockchain/head.go b/beacon-chain/blockchain/head.go
index 405311a0b..0cd9a95ee 100644
--- a/beacon-chain/blockchain/head.go
+++ b/beacon-chain/blockchain/head.go
@@ -64,6 +64,9 @@ func (s *Service) saveHead(ctx context.Context, headRoot [32]byte) error {
 	if err != nil {
 		return errors.Wrap(err, "could not retrieve head state in DB")
 	}
+	if headState == nil {
+		return errors.New("cannot save nil head state")
+	}
 
 	s.headLock.Lock()
 	defer s.headLock.Unlock()