core, dex, downloader: refactor governance

1 files changed, 60 insertions, 24 deletions

diff --git a/core/governance.go b/core/governance.go
index 261db1e75..dcc390eb4 100644
--- a/core/governance.go
+++ b/core/governance.go
@@ -1,6 +1,7 @@
 package core
 
 import (
+	"fmt"
 	"math/big"
 	"time"
 
@@ -8,29 +9,72 @@ import (
 	coreTypes "github.com/dexon-foundation/dexon-consensus/core/types"
 	dkgTypes "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
 
+	"github.com/dexon-foundation/dexon/core/state"
 	"github.com/dexon-foundation/dexon/core/vm"
 	"github.com/dexon-foundation/dexon/log"
 	"github.com/dexon-foundation/dexon/rlp"
 )
 
-type Governance struct {
+type GovernanceStateDB interface {
+	State() (*state.StateDB, error)
+	StateAt(height uint64) (*state.StateDB, error)
+}
+
+func NewGovernanceStateDB(bc *BlockChain) GovernanceStateDB {
+	return &governanceStateDB{bc: bc}
+}
+
+type governanceStateDB struct {
 	bc *BlockChain
 }
 
-func NewGovernance(bc *BlockChain) *Governance {
-	return &Governance{bc: bc}
+func (g *governanceStateDB) State() (*state.StateDB, error) {
+	return g.bc.State()
 }
 
-func (g *Governance) getHeadHelper() *vm.GovernanceStateHelper {
-	headState, err := g.bc.State()
+func (g *governanceStateDB) StateAt(height uint64) (*state.StateDB, error) {
+	header := g.bc.GetHeaderByNumber(height)
+	if header == nil {
+		return nil, fmt.Errorf("header at %d not exists", height)
+	}
+	return g.bc.StateAt(header.Root)
+}
+
+type Governance struct {
+	db GovernanceStateDB
+}
+
+func NewGovernance(db GovernanceStateDB) *Governance {
+	return &Governance{db: db}
+}
+
+func (g *Governance) GetHeadHelper() *vm.GovernanceStateHelper {
+	headState, err := g.db.State()
 	if err != nil {
-		log.Error("get head state fail", "err", err)
+		log.Error("Governance head state not ready", "err", err)
 		panic(err)
 	}
 	return &vm.GovernanceStateHelper{headState}
 }
 
-func (g *Governance) getConfigHelper(round uint64) *vm.GovernanceStateHelper {
+func (g *Governance) getHelperAtRound(round uint64) *vm.GovernanceStateHelper {
+	height := g.GetRoundHeight(round)
+
+	// Sanity check
+	if round != 0 && height == 0 {
+		log.Error("Governance state incorrect", "round", round, "got height", height)
+		panic("round != 0 but height == 0")
+	}
+
+	s, err := g.db.StateAt(height)
+	if err != nil {
+		log.Error("Governance state not ready", "round", round, "height", height, "err", err)
+		panic(err)
+	}
+	return &vm.GovernanceStateHelper{s}
+}
+
+func (g *Governance) GetConfigHelper(round uint64) *vm.GovernanceStateHelper {
 	if round < dexCore.ConfigRoundShift {
 		round = 0
 	} else {
@@ -39,21 +83,13 @@ func (g *Governance) getConfigHelper(round uint64) *vm.GovernanceStateHelper {
 	return g.getHelperAtRound(round)
 }
 
-func (g *Governance) getHelperAtRound(round uint64) *vm.GovernanceStateHelper {
-	headHelper := g.getHeadHelper()
-	height := headHelper.RoundHeight(big.NewInt(int64(round))).Uint64()
-	header := g.bc.GetHeaderByNumber(height)
-	s, err := g.bc.StateAt(header.Root)
-	if err != nil {
-		log.Error("get state fail", "err", err)
-		panic(err)
-	}
-	return &vm.GovernanceStateHelper{s}
+func (g *Governance) GetRoundHeight(round uint64) uint64 {
+	return g.GetHeadHelper().RoundHeight(big.NewInt(int64(round))).Uint64()
 }
 
 func (g *Governance) Configuration(round uint64) *coreTypes.Config {
-	helper := g.getConfigHelper(round)
-	c := helper.Configuration()
+	configHelper := g.GetConfigHelper(round)
+	c := configHelper.Configuration()
 	return &coreTypes.Config{
 		NumChains:        c.NumChains,
 		LambdaBA:         time.Duration(c.LambdaBA) * time.Millisecond,
@@ -68,7 +104,7 @@ func (g *Governance) Configuration(round uint64) *coreTypes.Config {
 }
 
 func (g *Governance) DKGComplaints(round uint64) []*dkgTypes.Complaint {
-	headHelper := g.getHeadHelper()
+	headHelper := g.GetHeadHelper()
 	var dkgComplaints []*dkgTypes.Complaint
 	for _, pk := range headHelper.DKGComplaints(big.NewInt(int64(round))) {
 		x := new(dkgTypes.Complaint)
@@ -81,7 +117,7 @@ func (g *Governance) DKGComplaints(round uint64) []*dkgTypes.Complaint {
 }
 
 func (g *Governance) DKGMasterPublicKeys(round uint64) []*dkgTypes.MasterPublicKey {
-	headHelper := g.getHeadHelper()
+	headHelper := g.GetHeadHelper()
 	var dkgMasterPKs []*dkgTypes.MasterPublicKey
 	for _, pk := range headHelper.DKGMasterPublicKeys(big.NewInt(int64(round))) {
 		x := new(dkgTypes.MasterPublicKey)
@@ -94,9 +130,9 @@ func (g *Governance) DKGMasterPublicKeys(round uint64) []*dkgTypes.MasterPublicK
 }
 
 func (g *Governance) IsDKGFinal(round uint64) bool {
-	headHelper := g.getHeadHelper()
-	configHelper := g.getConfigHelper(round)
-	threshold := 2*configHelper.DKGSetSize().Uint64()/3 + 1
+	headHelper := g.GetHeadHelper()
+	config := g.Configuration(round)
+	threshold := 2*uint64(config.DKGSetSize)/3 + 1
 	count := headHelper.DKGFinalizedsCount(big.NewInt(int64(round))).Uint64()
 	return count >= threshold
 }