From f2524a888d6133e7d729ee22031d198695e8ce5f Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Sun, 5 Apr 2020 12:33:35 -0700
Subject: [PATCH] Make validator snapshot on new validator too (#2733)

---
 core/blockchain.go | 21 +++++++++++----------
 core/offchain.go   |  7 ++++---
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/core/blockchain.go b/core/blockchain.go
index 120da215d..6f2935427 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -2430,7 +2430,7 @@ func (bc *BlockChain) UpdateValidatorVotingPower(
 // UpdateValidatorSnapshots updates the content snapshot of all validators
 // Note: this should only be called within the blockchain insert process.
 func (bc *BlockChain) UpdateValidatorSnapshots(
-	batch rawdb.DatabaseWriter, epoch *big.Int, state *state.DB,
+	batch rawdb.DatabaseWriter, epoch *big.Int, state *state.DB, newValidators []common.Address,
 ) error {
 	// Note this is reading the validator list from last block.
 	// It's fine since the new validators from this block is already snapshot when created.
@@ -2439,6 +2439,7 @@ func (bc *BlockChain) UpdateValidatorSnapshots(
 		return err
 	}
 
+	allValidators = append(allValidators, newValidators...)
 	// TODO: enable this once we allow validator to delete itself.
 	//err = bc.deleteValidatorSnapshots(allValidators)
 	//if err != nil {
@@ -2544,49 +2545,49 @@ func (bc *BlockChain) writeDelegationsByDelegator(
 func (bc *BlockChain) UpdateStakingMetaData(
 	batch rawdb.DatabaseWriter, txns staking.StakingTransactions,
 	state *state.DB, epoch *big.Int,
-) error {
+) (newValidators []common.Address, err error) {
 	newValidators, newDelegations, err := bc.prepareStakingMetaData(txns, state)
 	if err != nil {
 		utils.Logger().Warn().Msgf("oops, prepareStakingMetaData failed, err: %+v", err)
-		return err
+		return newValidators, err
 	}
 
 	if len(newValidators) > 0 {
 		list, err := bc.ReadValidatorList()
 		if err != nil {
-			return err
+			return newValidators, err
 		}
 
 		for _, addr := range newValidators {
 			newList, appended := utils.AppendIfMissing(list, addr)
 			if !appended {
-				return errValidatorExist
+				return newValidators, errValidatorExist
 			}
 			list = newList
 
 			// Update validator snapshot for the new validator
 			validator, err := state.ValidatorWrapper(addr)
 			if err != nil {
-				return err
+				return newValidators, err
 			}
 
 			if err := rawdb.WriteValidatorSnapshot(batch, validator, epoch); err != nil {
-				return err
+				return newValidators, err
 			}
 		}
 
 		// Update validator list
 		if err = bc.WriteValidatorList(batch, list); err != nil {
-			return err
+			return newValidators, err
 		}
 	}
 
 	for addr, delegations := range newDelegations {
 		if err := bc.writeDelegationsByDelegator(batch, addr, delegations); err != nil {
-			return err
+			return newValidators, err
 		}
 	}
-	return nil
+	return newValidators, nil
 }
 
 // prepareStakingMetaData prepare the updates of validator's
diff --git a/core/offchain.go b/core/offchain.go
index d5e9195bd..e72a0a3a8 100644
--- a/core/offchain.go
+++ b/core/offchain.go
@@ -93,9 +93,10 @@ func (bc *BlockChain) CommitOffChainData(
 	//}
 
 	// Do bookkeeping for new staking txns
-	if err := bc.UpdateStakingMetaData(
+	newVals, err := bc.UpdateStakingMetaData(
 		batch, block.StakingTransactions(), state, epoch,
-	); err != nil {
+	)
+	if err != nil {
 		utils.Logger().Err(err).Msg("UpdateStakingMetaData failed")
 		return NonStatTy, err
 	}
@@ -129,7 +130,7 @@ func (bc *BlockChain) CommitOffChainData(
 	if isBeaconChain && shard.Schedule.IsLastBlock(header.Number().Uint64()+1) {
 		// Update snapshots for all validators
 		epoch := new(big.Int).Add(header.Epoch(), common.Big1)
-		if err := bc.UpdateValidatorSnapshots(batch, epoch, state); err != nil {
+		if err := bc.UpdateValidatorSnapshots(batch, epoch, state, newVals); err != nil {
 			return NonStatTy, err
 		}
 	}