core, dex, downloader: refactor governance

6 files changed, 109 insertions, 246 deletions

diff --git a/core/blockchain.go b/core/blockchain.go
index d9401731c..3d7171cd7 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -234,7 +234,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par
 		}
 	}
 
-	gov := NewGovernance(bc)
+	gov := NewGovernance(NewGovernanceStateDB(bc))
 	bc.verifierCache = dexCore.NewTSigVerifierCache(gov, 5)
 
 	// Take ownership of this particular state
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
 }
diff --git a/dex/downloader/governance.go b/dex/downloader/governance.go
index 40233f1a8..38281b4a8 100644
--- a/dex/downloader/governance.go
+++ b/dex/downloader/governance.go
@@ -1,27 +1,22 @@
 package downloader
 
 import (
-	"math/big"
+	"fmt"
 	"sync"
-	"time"
-
-	dexCore "github.com/dexon-foundation/dexon-consensus/core"
-	coreTypes "github.com/dexon-foundation/dexon-consensus/core/types"
-	dkgTypes "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
 
 	"github.com/dexon-foundation/dexon/common"
+	"github.com/dexon-foundation/dexon/core"
 	"github.com/dexon-foundation/dexon/core/state"
 	"github.com/dexon-foundation/dexon/core/types"
-	"github.com/dexon-foundation/dexon/core/vm"
 	"github.com/dexon-foundation/dexon/crypto"
 	"github.com/dexon-foundation/dexon/ethdb"
 	"github.com/dexon-foundation/dexon/log"
-	"github.com/dexon-foundation/dexon/rlp"
 	"github.com/dexon-foundation/dexon/trie"
 )
 
-// This is a goverance for fast sync
-type governance struct {
+// governanceDB is backed by memory db for fast sync.
+// it implements core.GovernanceStateDB
+type governanceStateDB struct {
 	db          ethdb.Database
 	headRoot    common.Hash
 	headHeight  uint64
@@ -29,66 +24,32 @@ type governance struct {
 	mu          sync.Mutex
 }
 
-func newGovernance(headState *types.GovState) *governance {
-	db := ethdb.NewMemDatabase()
-	g := &governance{
-		db:          db,
-		headRoot:    headState.Root,
-		headHeight:  headState.Number.Uint64(),
-		height2Root: make(map[uint64]common.Hash),
-	}
-	g.StoreState(headState)
-
-	statedb, err := state.New(headState.Root, state.NewDatabase(g.db))
-	if statedb == nil || err != nil {
-		log.Error("New governance fail", "statedb == nil", statedb == nil, "err", err)
-	}
-	return g
-}
-
-func (g *governance) getHeadHelper() *vm.GovernanceStateHelper {
-	return g.getHelper(g.headRoot)
+func (g *governanceStateDB) State() (*state.StateDB, error) {
+	return state.New(g.headRoot, state.NewDatabase(g.db))
 }
 
-func (g *governance) getHelperByRound(round uint64) *vm.GovernanceStateHelper {
-	height := g.GetRoundHeight(round)
+func (g *governanceStateDB) StateAt(height uint64) (*state.StateDB, error) {
 	root, exists := g.height2Root[height]
 	if !exists {
-		log.Debug("Gov get helper by round", "round", round, "exists", exists)
-		return nil
+		return nil, fmt.Errorf("Governance state not ready, height: %d", height)
 	}
-	return g.getHelper(root)
+	return state.New(root, state.NewDatabase(g.db))
 }
 
-func (g *governance) getHelper(root common.Hash) *vm.GovernanceStateHelper {
-	statedb, err := state.New(root, state.NewDatabase(g.db))
-	if statedb == nil || err != nil {
-		return nil
-	}
-	return &vm.GovernanceStateHelper{statedb}
-}
-
-func (g *governance) GetRoundHeight(round uint64) uint64 {
-	var h uint64
-	if helper := g.getHeadHelper(); helper != nil {
-		h = helper.RoundHeight(big.NewInt(int64(round))).Uint64()
-	}
-	return h
-}
-
-func (g *governance) StoreState(s *types.GovState) {
+func (g *governanceStateDB) StoreState(s *types.GovState) {
 	g.mu.Lock()
 	defer g.mu.Unlock()
+	log.Debug("Store state", "height", s.Number.Uint64())
 
-	// store the hight -> root mapping
+	// Store the height -> root mapping.
 	g.height2Root[s.Number.Uint64()] = s.Root
 
-	// store the account
+	// Store the account.
 	for _, node := range s.Proof {
 		g.db.Put(crypto.Keccak256(node), node)
 	}
 
-	// store the storage
+	// Store the storage.
 	triedb := trie.NewDatabase(g.db)
 	t, err := trie.New(common.Hash{}, triedb)
 	if err != nil {
@@ -101,70 +62,33 @@ func (g *governance) StoreState(s *types.GovState) {
 	triedb.Commit(t.Hash(), false)
 
 	if s.Number.Uint64() > g.headHeight {
-		log.Debug("Gov head root changed", "number", s.Number.Uint64())
+		log.Debug("Governance head root changed", "number", s.Number.Uint64())
 		g.headRoot = s.Root
 		g.headHeight = s.Number.Uint64()
 	}
 }
 
-// Return the genesis configuration if round == 0.
-func (g *governance) Configuration(round uint64) *coreTypes.Config {
-	if round < dexCore.ConfigRoundShift {
-		round = 0
-	} else {
-		round -= dexCore.ConfigRoundShift
-	}
-	helper := g.getHelperByRound(round)
-	if helper == nil {
-		log.Warn("Get config helper fail", "round - round shift", round)
-		return nil
-	}
-	c := helper.Configuration()
-	return &coreTypes.Config{
-		NumChains:        c.NumChains,
-		LambdaBA:         time.Duration(c.LambdaBA) * time.Millisecond,
-		LambdaDKG:        time.Duration(c.LambdaDKG) * time.Millisecond,
-		K:                int(c.K),
-		PhiRatio:         c.PhiRatio,
-		NotarySetSize:    c.NotarySetSize,
-		DKGSetSize:       c.DKGSetSize,
-		RoundInterval:    time.Duration(c.RoundInterval) * time.Millisecond,
-		MinBlockInterval: time.Duration(c.MinBlockInterval) * time.Millisecond,
-	}
+// This is a governance for fast sync
+type governance struct {
+	*core.Governance
+	db *governanceStateDB
 }
 
-// DKGComplaints gets all the DKGComplaints of round.
-func (g *governance) DKGComplaints(round uint64) []*dkgTypes.Complaint {
-	helper := g.getHeadHelper()
-	var dkgComplaints []*dkgTypes.Complaint
-	for _, pk := range helper.DKGComplaints(big.NewInt(int64(round))) {
-		x := new(dkgTypes.Complaint)
-		if err := rlp.DecodeBytes(pk, x); err != nil {
-			panic(err)
-		}
-		dkgComplaints = append(dkgComplaints, x)
+func newGovernance(headState *types.GovState) *governance {
+	db := ethdb.NewMemDatabase()
+	govStateDB := &governanceStateDB{
+		db:          db,
+		headRoot:    headState.Root,
+		headHeight:  headState.Number.Uint64(),
+		height2Root: make(map[uint64]common.Hash),
 	}
-	return dkgComplaints
-}
-
-// DKGMasterPublicKeys gets all the DKGMasterPublicKey of round.
-func (g *governance) DKGMasterPublicKeys(round uint64) []*dkgTypes.MasterPublicKey {
-	helper := g.getHeadHelper()
-	var dkgMasterPKs []*dkgTypes.MasterPublicKey
-	for _, pk := range helper.DKGMasterPublicKeys(big.NewInt(int64(round))) {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(pk, x); err != nil {
-			panic(err)
-		}
-		dkgMasterPKs = append(dkgMasterPKs, x)
+	govStateDB.StoreState(headState)
+	return &governance{
+		Governance: core.NewGovernance(govStateDB),
+		db:         govStateDB,
 	}
-	return dkgMasterPKs
 }
 
-// IsDKGFinal checks if DKG is final.
-func (g *governance) IsDKGFinal(round uint64) bool {
-	helper := g.getHeadHelper()
-	threshold := 2*uint64(g.Configuration(round).DKGSetSize)/3 + 1
-	count := helper.DKGFinalizedsCount(big.NewInt(int64(round))).Uint64()
-	return count >= threshold
+func (g *governance) StoreState(s *types.GovState) {
+	g.db.StoreState(s)
 }
diff --git a/dex/governance.go b/dex/governance.go
index 01f39b68c..aeb6d3fb0 100644
--- a/dex/governance.go
+++ b/dex/governance.go
@@ -5,26 +5,26 @@ import (
 	"crypto/ecdsa"
 	"encoding/hex"
 	"math/big"
-	"time"
 
 	coreCommon "github.com/dexon-foundation/dexon-consensus/common"
 	dexCore "github.com/dexon-foundation/dexon-consensus/core"
 	coreCrypto "github.com/dexon-foundation/dexon-consensus/core/crypto"
 	coreEcdsa "github.com/dexon-foundation/dexon-consensus/core/crypto/ecdsa"
-	coreTypes "github.com/dexon-foundation/dexon-consensus/core/types"
 	dkgTypes "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
 
 	"github.com/dexon-foundation/dexon/common"
+	"github.com/dexon-foundation/dexon/core"
 	"github.com/dexon-foundation/dexon/core/types"
 	"github.com/dexon-foundation/dexon/core/vm"
 	"github.com/dexon-foundation/dexon/crypto"
 	"github.com/dexon-foundation/dexon/log"
 	"github.com/dexon-foundation/dexon/params"
 	"github.com/dexon-foundation/dexon/rlp"
-	"github.com/dexon-foundation/dexon/rpc"
 )
 
 type DexconGovernance struct {
+	*core.Governance
+
 	b            *DexAPIBackend
 	chainConfig  *params.ChainConfig
 	privateKey   *ecdsa.PrivateKey
@@ -32,11 +32,13 @@ type DexconGovernance struct {
 	nodeSetCache *dexCore.NodeSetCache
 }
 
-// NewDexconGovernance retruns a governance implementation of the DEXON
+// NewDexconGovernance returns a governance implementation of the DEXON
 // consensus governance interface.
 func NewDexconGovernance(backend *DexAPIBackend, chainConfig *params.ChainConfig,
 	privKey *ecdsa.PrivateKey) *DexconGovernance {
 	g := &DexconGovernance{
+		Governance: core.NewGovernance(
+			core.NewGovernanceStateDB(backend.dex.BlockChain())),
 		b:           backend,
 		chainConfig: chainConfig,
 		privateKey:  privKey,
@@ -46,67 +48,9 @@ func NewDexconGovernance(backend *DexAPIBackend, chainConfig *params.ChainConfig
 	return g
 }
 
-func (d *DexconGovernance) GetRoundHeight(ctx context.Context, round uint64) (uint64, error) {
-	state, _, err := d.b.StateAndHeaderByNumber(ctx, rpc.LatestBlockNumber)
-	if state == nil || err != nil {
-		return 0, err
-	}
-	s := vm.GovernanceStateHelper{state}
-	return s.RoundHeight(big.NewInt(int64(round))).Uint64(), nil
-}
-
-func (d *DexconGovernance) getGovState() *vm.GovernanceStateHelper {
-	ctx := context.Background()
-	state, _, err := d.b.StateAndHeaderByNumber(ctx, rpc.LatestBlockNumber)
-	if state == nil || err != nil {
-		return nil
-	}
-
-	return &vm.GovernanceStateHelper{state}
-}
-
-func (d *DexconGovernance) getGovStateAtRound(round uint64) *vm.GovernanceStateHelper {
-	if round < dexCore.ConfigRoundShift {
-		round = 0
-	} else {
-		round -= dexCore.ConfigRoundShift
-	}
-	ctx := context.Background()
-	blockHeight, err := d.GetRoundHeight(ctx, round)
-	if err != nil {
-		return nil
-	}
-
-	state, _, err := d.b.StateAndHeaderByNumber(ctx, rpc.BlockNumber(blockHeight))
-	if state == nil || err != nil {
-		log.Error("Failed to retrieve governance state", "round", round, "height", blockHeight)
-		return nil
-	}
-	return &vm.GovernanceStateHelper{state}
-}
-
 // DexconConfiguration return raw config in state.
 func (d *DexconGovernance) DexconConfiguration(round uint64) *params.DexconConfig {
-	s := d.getGovStateAtRound(round)
-	return s.Configuration()
-}
-
-// Configuration returns the system configuration for consensus core to use.
-func (d *DexconGovernance) Configuration(round uint64) *coreTypes.Config {
-	s := d.getGovStateAtRound(round)
-	c := s.Configuration()
-
-	return &coreTypes.Config{
-		NumChains:        c.NumChains,
-		LambdaBA:         time.Duration(c.LambdaBA) * time.Millisecond,
-		LambdaDKG:        time.Duration(c.LambdaDKG) * time.Millisecond,
-		K:                int(c.K),
-		PhiRatio:         c.PhiRatio,
-		NotarySetSize:    c.NotarySetSize,
-		DKGSetSize:       c.DKGSetSize,
-		RoundInterval:    time.Duration(c.RoundInterval) * time.Millisecond,
-		MinBlockInterval: time.Duration(c.MinBlockInterval) * time.Millisecond,
-	}
+	return d.GetConfigHelper(round).Configuration()
 }
 
 func (d *DexconGovernance) sendGovTx(ctx context.Context, data []byte) error {
@@ -146,12 +90,12 @@ func (d *DexconGovernance) sendGovTx(ctx context.Context, data []byte) error {
 
 // CRS returns the CRS for a given round.
 func (d *DexconGovernance) CRS(round uint64) coreCommon.Hash {
-	s := d.getGovState()
+	s := d.GetHeadHelper()
 	return coreCommon.Hash(s.CRS(big.NewInt(int64(round))))
 }
 
 func (d *DexconGovernance) LenCRS() uint64 {
-	s := d.getGovState()
+	s := d.GetHeadHelper()
 	return s.LenCRS().Uint64()
 }
 
@@ -174,7 +118,7 @@ func (d *DexconGovernance) ProposeCRS(round uint64, signedCRS []byte) {
 
 // NodeSet returns the current node set.
 func (d *DexconGovernance) NodeSet(round uint64) []coreCrypto.PublicKey {
-	s := d.getGovStateAtRound(round)
+	s := d.GetConfigHelper(round)
 	var pks []coreCrypto.PublicKey
 
 	for _, n := range s.QualifiedNodes() {
@@ -228,20 +172,6 @@ func (d *DexconGovernance) AddDKGComplaint(round uint64, complaint *dkgTypes.Com
 	}
 }
 
-// DKGComplaints gets all the DKGComplaints of round.
-func (d *DexconGovernance) DKGComplaints(round uint64) []*dkgTypes.Complaint {
-	s := d.getGovState()
-	var dkgComplaints []*dkgTypes.Complaint
-	for _, pk := range s.DKGComplaints(big.NewInt(int64(round))) {
-		x := new(dkgTypes.Complaint)
-		if err := rlp.DecodeBytes(pk, x); err != nil {
-			panic(err)
-		}
-		dkgComplaints = append(dkgComplaints, x)
-	}
-	return dkgComplaints
-}
-
 // AddDKGMasterPublicKey adds a DKGMasterPublicKey.
 func (d *DexconGovernance) AddDKGMasterPublicKey(round uint64, masterPublicKey *dkgTypes.MasterPublicKey) {
 	method := vm.GovernanceContractName2Method["addDKGMasterPublicKey"]
@@ -265,20 +195,6 @@ func (d *DexconGovernance) AddDKGMasterPublicKey(round uint64, masterPublicKey *
 	}
 }
 
-// DKGMasterPublicKeys gets all the DKGMasterPublicKey of round.
-func (d *DexconGovernance) DKGMasterPublicKeys(round uint64) []*dkgTypes.MasterPublicKey {
-	s := d.getGovState()
-	var dkgMasterPKs []*dkgTypes.MasterPublicKey
-	for _, pk := range s.DKGMasterPublicKeys(big.NewInt(int64(round))) {
-		x := new(dkgTypes.MasterPublicKey)
-		if err := rlp.DecodeBytes(pk, x); err != nil {
-			panic(err)
-		}
-		dkgMasterPKs = append(dkgMasterPKs, x)
-	}
-	return dkgMasterPKs
-}
-
 // AddDKGFinalize adds a DKG finalize message.
 func (d *DexconGovernance) AddDKGFinalize(round uint64, final *dkgTypes.Finalize) {
 	method := vm.GovernanceContractName2Method["addDKGFinalize"]
@@ -302,14 +218,6 @@ func (d *DexconGovernance) AddDKGFinalize(round uint64, final *dkgTypes.Finalize
 	}
 }
 
-// IsDKGFinal checks if DKG is final.
-func (d *DexconGovernance) IsDKGFinal(round uint64) bool {
-	s := d.getGovState()
-	threshold := 2*s.DKGSetSize().Uint64()/3 + 1
-	count := s.DKGFinalizedsCount(big.NewInt(int64(round))).Uint64()
-	return count >= threshold
-}
-
 func (d *DexconGovernance) GetNumChains(round uint64) uint32 {
 	return d.Configuration(round).NumChains
 }
diff --git a/dex/handler.go b/dex/handler.go
index a86a68329..b4499ae5a 100644
--- a/dex/handler.go
+++ b/dex/handler.go
@@ -34,7 +34,6 @@
 package dex
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -495,7 +494,6 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 			}
 		}
 
-		ctx := context.Background()
 		if query.WithGov && len(headers) > 0 {
 			last := headers[len(headers)-1]
 			currentBlock := pm.blockchain.CurrentBlock()
@@ -513,11 +511,7 @@ func (pm *ProtocolManager) handleMsg(p *peer) error {
 				if r == 0 {
 					continue
 				}
-				h, err := pm.gov.GetRoundHeight(ctx, r)
-				if err != nil {
-					log.Warn("Get round height fail", "err", err)
-					return p.SendBlockHeaders([]*types.HeaderWithGovState{})
-				}
+				h := pm.gov.GetRoundHeight(r)
 				log.Trace("#Snapshot height", "height", h)
 				if h == 0 {
 					h = height
@@ -1181,7 +1175,9 @@ func (pm *ProtocolManager) peerSetLoop() {
 			}
 			if newRound == round+1 {
 				pm.peers.BuildConnection(newRound)
-				pm.peers.ForgetConnection(round - 1)
+				if round >= 1 {
+					pm.peers.ForgetConnection(round - 1)
+				}
 			} else {
 				// just forget all network connection and rebuild.
 				pm.peers.ForgetConnection(round)
diff --git a/dex/protocol.go b/dex/protocol.go
index 0c271248e..21168e917 100644
--- a/dex/protocol.go
+++ b/dex/protocol.go
@@ -34,7 +34,6 @@
 package dex
 
 import (
-	"context"
 	"crypto/ecdsa"
 	"fmt"
 	"io"
@@ -144,7 +143,7 @@ type txPool interface {
 }
 
 type governance interface {
-	GetRoundHeight(context.Context, uint64) (uint64, error)
+	GetRoundHeight(uint64) uint64
 
 	GetNumChains(uint64) uint32