From 5eb3684d8973493929206a10c68ec17fe2d872ed Mon Sep 17 00:00:00 2001
From: chaosma <chao@harmony.one>
Date: Wed, 13 Feb 2019 15:02:39 -0800
Subject: [PATCH] fix shard state update bug (#439)

---
 core/blockchain.go   | 4 ++++
 node/node_handler.go | 2 --
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 90decdb11..86075f52c 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -1244,11 +1244,15 @@ func (bc *BlockChain) insertChain(chain types.Blocks) (int, []interface{}, []*ty
 
 		cache, _ := bc.stateCache.TrieDB().Size()
 		stats.report(chain, i, cache)
+
+		// only insert new shardstate when block is epoch block
+		bc.InsertNewShardState(block)
 	}
 	// Append a single chain head event if we've progressed the chain
 	if lastCanon != nil && bc.CurrentBlock().Hash() == lastCanon.Hash() {
 		events = append(events, ChainHeadEvent{lastCanon})
 	}
+
 	return 0, events, coalescedLogs, nil
 }
 
diff --git a/node/node_handler.go b/node/node_handler.go
index 3c8da7610..833aa21c5 100644
--- a/node/node_handler.go
+++ b/node/node_handler.go
@@ -277,8 +277,6 @@ func (node *Node) AddNewBlock(newBlock *types.Block) {
 	} else {
 		utils.GetLogInstance().Info("adding new block to blockchain", "blockNum", blockNum)
 	}
-	// only insert new shardstate when newBlock is epoch block
-	node.blockchain.InsertNewShardState(newBlock)
 }
 
 func (node *Node) pingMessageHandler(msgPayload []byte) int {