dex: gov: using dex-consensus-core NodeSetCache

5 files changed, 71 insertions, 64 deletions

diff --git a/dex/governance.go b/dex/governance.go
index 32c2d79af..22452dea3 100644
--- a/dex/governance.go
+++ b/dex/governance.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	coreCommon "github.com/dexon-foundation/dexon-consensus-core/common"
+	dexCore "github.com/dexon-foundation/dexon-consensus-core/core"
 	coreCrypto "github.com/dexon-foundation/dexon-consensus-core/core/crypto"
 	coreEcdsa "github.com/dexon-foundation/dexon-consensus-core/core/crypto/ecdsa"
 	coreTypes "github.com/dexon-foundation/dexon-consensus-core/core/types"
@@ -23,22 +24,25 @@ import (
 )
 
 type DexconGovernance struct {
-	b           *DexAPIBackend
-	chainConfig *params.ChainConfig
-	privateKey  *ecdsa.PrivateKey
-	address     common.Address
+	b            *DexAPIBackend
+	chainConfig  *params.ChainConfig
+	privateKey   *ecdsa.PrivateKey
+	address      common.Address
+	nodeSetCache *dexCore.NodeSetCache
 }
 
 // NewDexconGovernance retruns a governance implementation of the DEXON
 // consensus governance interface.
 func NewDexconGovernance(backend *DexAPIBackend, chainConfig *params.ChainConfig,
 	privKey *ecdsa.PrivateKey) *DexconGovernance {
-	return &DexconGovernance{
+	g := &DexconGovernance{
 		b:           backend,
 		chainConfig: chainConfig,
 		privateKey:  privKey,
 		address:     crypto.PubkeyToAddress(privKey.PublicKey),
 	}
+	g.nodeSetCache = dexCore.NewNodeSetCache(g)
+	return g
 }
 
 func (d *DexconGovernance) getRoundHeight(ctx context.Context, round uint64) (uint64, error) {
@@ -283,60 +287,41 @@ func (d *DexconGovernance) GetNumChains(round uint64) uint32 {
 	return d.Configuration(round).NumChains
 }
 
-func (d *DexconGovernance) NotarySet(chainID uint32, round uint64) map[string]struct{} {
-	id2Key := map[coreTypes.NodeID]coreCrypto.PublicKey{}
-
-	nodeSet := coreTypes.NewNodeSet()
-	for _, key := range d.NodeSet(round) {
-		id := coreTypes.NewNodeID(key)
-		id2Key[id] = key
-		nodeSet.Add(id)
+func (d *DexconGovernance) NotarySet(
+	round uint64, chainID uint32) (map[string]struct{}, error) {
+	notarySet, err := d.nodeSetCache.GetNotarySet(round, chainID)
+	if err != nil {
+		return nil, err
 	}
 
-	cfg := d.Configuration(round)
-	crs := d.CRS(round)
-
-	notarySet := nodeSet.GetSubSet(
-		int(cfg.NotarySetSize), coreTypes.NewNotarySetTarget(crs, chainID))
-
-	r := map[string]struct{}{}
+	r := make(map[string]struct{}, len(notarySet))
 	for id := range notarySet {
-		compressed := id2Key[id]
-		// 33 bytes pubkey to 65 bytes pubkey
-		key, err := crypto.DecompressPubkey(compressed.Bytes())
-		if err != nil {
-			continue
+		if key, exists := d.nodeSetCache.GetPublicKey(id); exists {
+			uncompressedKey, err := crypto.DecompressPubkey(key.Bytes())
+			if err != nil {
+				log.Error("decompress key fail", "err", err)
+			}
+			r[discover.PubkeyID(uncompressedKey).String()] = struct{}{}
 		}
-		r[discover.PubkeyID(key).String()] = struct{}{}
 	}
-	return r
+	return r, nil
 }
 
-func (d *DexconGovernance) DKGSet(round uint64) map[string]struct{} {
-	id2Key := map[coreTypes.NodeID]coreCrypto.PublicKey{}
-
-	nodeSet := coreTypes.NewNodeSet()
-	for _, key := range d.NodeSet(round) {
-		id := coreTypes.NewNodeID(key)
-		id2Key[id] = key
-		nodeSet.Add(id)
+func (d *DexconGovernance) DKGSet(round uint64) (map[string]struct{}, error) {
+	dkgSet, err := d.nodeSetCache.GetDKGSet(round)
+	if err != nil {
+		return nil, err
 	}
 
-	cfg := d.Configuration(round)
-	crs := d.CRS(round)
-
-	dkgSet := nodeSet.GetSubSet(
-		int(cfg.DKGSetSize), coreTypes.NewDKGSetTarget(crs))
-
-	r := map[string]struct{}{}
+	r := make(map[string]struct{}, len(dkgSet))
 	for id := range dkgSet {
-		compressed := id2Key[id]
-		// 33 bytes pubkey to 65 bytes pubkey
-		key, err := crypto.DecompressPubkey(compressed.Bytes())
-		if err != nil {
-			continue
+		if key, exists := d.nodeSetCache.GetPublicKey(id); exists {
+			uncompressedKey, err := crypto.DecompressPubkey(key.Bytes())
+			if err != nil {
+				log.Error("decompress key fail", "err", err)
+			}
+			r[discover.PubkeyID(uncompressedKey).String()] = struct{}{}
 		}
-		r[discover.PubkeyID(key).String()] = struct{}{}
 	}
-	return r
+	return r, nil
 }
diff --git a/dex/helper_test.go b/dex/helper_test.go
index 21803ed88..09d15d42f 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -183,19 +183,20 @@ func newTestTransaction(from *ecdsa.PrivateKey, nonce uint64, datasize int) *typ
 // testGovernance is a fake, helper governance for testing purposes
 type testGovernance struct {
 	numChainsFunc func(uint64) uint32
-	notarySetFunc func(uint32, uint64) map[string]struct{}
-	dkgSetFunc    func(uint64) map[string]struct{}
+	notarySetFunc func(uint64, uint32) (map[string]struct{}, error)
+	dkgSetFunc    func(uint64) (map[string]struct{}, error)
 }
 
 func (g *testGovernance) GetNumChains(round uint64) uint32 {
 	return g.numChainsFunc(round)
 }
 
-func (g *testGovernance) NotarySet(chainID uint32, round uint64) map[string]struct{} {
-	return g.notarySetFunc(chainID, round)
+func (g *testGovernance) NotarySet(
+	round uint64, chainID uint32) (map[string]struct{}, error) {
+	return g.notarySetFunc(round, chainID)
 }
 
-func (g *testGovernance) DKGSet(round uint64) map[string]struct{} {
+func (g *testGovernance) DKGSet(round uint64) (map[string]struct{}, error) {
 	return g.dkgSetFunc(round)
 }
 
diff --git a/dex/peer.go b/dex/peer.go
index 44f61354c..342d0f033 100644
--- a/dex/peer.go
+++ b/dex/peer.go
@@ -27,6 +27,7 @@ import (
 	coreTypes "github.com/dexon-foundation/dexon-consensus-core/core/types"
 	"github.com/dexon-foundation/dexon/common"
 	"github.com/dexon-foundation/dexon/core/types"
+	"github.com/dexon-foundation/dexon/log"
 	"github.com/dexon-foundation/dexon/p2p"
 	"github.com/dexon-foundation/dexon/p2p/discover"
 	"github.com/dexon-foundation/dexon/rlp"
@@ -786,7 +787,12 @@ func (ps *peerSet) BuildNotaryConn(round uint64) {
 
 	selfID := ps.srvr.Self().ID.String()
 	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
-		s := ps.gov.NotarySet(chainID, round)
+		s, err := ps.gov.NotarySet(round, chainID)
+		if err != nil {
+			log.Error("get notary set fail",
+				"round", round, "chain id", chainID, "err", err)
+			continue
+		}
 
 		// not in notary set, add group
 		if _, ok := s[selfID]; !ok {
@@ -826,7 +832,12 @@ func (ps *peerSet) ForgetNotaryConn(round uint64) {
 func (ps *peerSet) forgetNotaryConn(round uint64) {
 	selfID := ps.srvr.Self().ID.String()
 	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
-		s := ps.gov.NotarySet(chainID, round)
+		s, err := ps.gov.NotarySet(round, chainID)
+		if err != nil {
+			log.Error("get notary set fail",
+				"round", round, "chain id", chainID, "err", err)
+			continue
+		}
 		if _, ok := s[selfID]; !ok {
 			ps.srvr.RemoveGroup(notarySetName(chainID, round))
 			continue
@@ -852,7 +863,12 @@ func (ps *peerSet) BuildDKGConn(round uint64) {
 	ps.lock.Lock()
 	defer ps.lock.Unlock()
 	selfID := ps.srvr.Self().ID.String()
-	s := ps.gov.DKGSet(round)
+	s, err := ps.gov.DKGSet(round)
+	if err != nil {
+		log.Error("get dkg set fail", "round", round)
+		return
+	}
+
 	if _, ok := s[selfID]; !ok {
 		return
 	}
@@ -882,7 +898,11 @@ func (ps *peerSet) ForgetDKGConn(round uint64) {
 
 func (ps *peerSet) forgetDKGConn(round uint64) {
 	selfID := ps.srvr.Self().ID.String()
-	s := ps.gov.DKGSet(round)
+	s, err := ps.gov.DKGSet(round)
+	if err != nil {
+		log.Error("get dkg set fail", "round", round)
+		return
+	}
 	if _, ok := s[selfID]; !ok {
 		return
 	}
diff --git a/dex/peer_test.go b/dex/peer_test.go
index 74989258a..c38848b66 100644
--- a/dex/peer_test.go
+++ b/dex/peer_test.go
@@ -42,13 +42,14 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		[]enode.ID{nodeID(0), nodeID(2), nodeID(6)},
 	}
 
-	gov.notarySetFunc = func(cid uint32, round uint64) map[string]struct{} {
+	gov.notarySetFunc = func(
+		round uint64, cid uint32) (map[string]struct{}, error) {
 		m := map[uint64][][]enode.ID{
 			10: round10,
 			11: round11,
 			12: round12,
 		}
-		return newTestNodeSet(m[round][cid])
+		return newTestNodeSet(m[round][cid]), nil
 	}
 
 	ps := newPeerSet(gov, server, table)
@@ -325,13 +326,13 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 
 	gov := &testGovernance{}
 
-	gov.dkgSetFunc = func(round uint64) map[string]struct{} {
+	gov.dkgSetFunc = func(round uint64) (map[string]struct{}, error) {
 		m := map[uint64][]enode.ID{
 			10: []enode.ID{nodeID(0), nodeID(1), nodeID(2)},
 			11: []enode.ID{nodeID(1), nodeID(2), nodeID(5)},
 			12: []enode.ID{nodeID(0), nodeID(3), nodeID(5)},
 		}
-		return newTestNodeSet(m[round])
+		return newTestNodeSet(m[round]), nil
 	}
 
 	ps := newPeerSet(gov, server, table)
diff --git a/dex/protocol.go b/dex/protocol.go
index 1907d18c8..8e1db583d 100644
--- a/dex/protocol.go
+++ b/dex/protocol.go
@@ -127,9 +127,9 @@ type txPool interface {
 type governance interface {
 	GetNumChains(uint64) uint32
 
-	NotarySet(uint32, uint64) map[string]struct{}
+	NotarySet(uint64, uint32) (map[string]struct{}, error)
 
-	DKGSet(uint64) map[string]struct{}
+	DKGSet(uint64) (map[string]struct{}, error)
 }
 
 type p2pServer interface {