dex: let notary nodes connect to some dkg nodes

BuildConnection builds notary and dkg connections together

1 files changed, 178 insertions, 38 deletions

diff --git a/dex/peer.go b/dex/peer.go
index 1f2b9518d..1279a190b 100644
--- a/dex/peer.go
+++ b/dex/peer.go
@@ -25,11 +25,13 @@ import (
 
 	mapset "github.com/deckarep/golang-set"
 	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/p2p/enode"
+	"github.com/dexon-foundation/dexon/p2p/enr"
 	"github.com/dexon-foundation/dexon/rlp"
 )
 
@@ -100,7 +102,6 @@ type peer struct {
 	rw p2p.MsgReadWriter
 
 	version int // Protocol version negotiated
-	labels  mapset.Set
 
 	head common.Hash
 	td   *big.Int
@@ -128,7 +129,6 @@ func newPeer(version int, p *p2p.Peer, rw p2p.MsgReadWriter) *peer {
 		Peer:                      p,
 		rw:                        rw,
 		version:                   version,
-		labels:                    mapset.NewSet(),
 		id:                        p.ID().String(),
 		knownTxs:                  mapset.NewSet(),
 		knownMetas:                mapset.NewSet(),
@@ -218,14 +218,6 @@ func (p *peer) close() {
 	close(p.term)
 }
 
-func (p *peer) addLabel(label peerLabel) {
-	p.labels.Add(label)
-}
-
-func (p *peer) removeLabel(label peerLabel) {
-	p.labels.Remove(label)
-}
-
 // Info gathers and returns a collection of metadata known about a peer.
 func (p *peer) Info() *PeerInfo {
 	hash, td := p.Head()
@@ -593,6 +585,7 @@ type peerSet struct {
 	gov           governance
 	peer2Labels   map[string]map[peerLabel]struct{}
 	label2Peers   map[peerLabel]map[string]struct{}
+	history       map[uint64]struct{}
 	notaryHistory map[uint64]struct{}
 	dkgHistory    map[uint64]struct{}
 }
@@ -606,6 +599,7 @@ func newPeerSet(gov governance, srvr p2pServer, tab *nodeTable) *peerSet {
 		tab:           tab,
 		peer2Labels:   make(map[string]map[peerLabel]struct{}),
 		label2Peers:   make(map[peerLabel]map[string]struct{}),
+		history:       make(map[uint64]struct{}),
 		notaryHistory: make(map[uint64]struct{}),
 		dkgHistory:    make(map[uint64]struct{}),
 	}
@@ -692,6 +686,18 @@ func (ps *peerSet) PeersWithoutTx(hash common.Hash) []*peer {
 	return list
 }
 
+func (ps *peerSet) PeersWithLabel(label peerLabel) []*peer {
+	ps.lock.RLock()
+	defer ps.lock.RUnlock()
+	list := make([]*peer, 0, len(ps.label2Peers[label]))
+	for id := range ps.label2Peers[label] {
+		if p, ok := ps.peers[id]; ok {
+			list = append(list, p)
+		}
+	}
+	return list
+}
+
 func (ps *peerSet) PeersWithoutVote(hash common.Hash, label peerLabel) []*peer {
 	ps.lock.RLock()
 	defer ps.lock.RUnlock()
@@ -775,6 +781,140 @@ func (ps *peerSet) Close() {
 	ps.closed = true
 }
 
+func (ps *peerSet) BuildConnection(round uint64) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+	defer ps.dumpPeerLabel(fmt.Sprintf("BuildConnection: %d", round))
+
+	ps.history[round] = struct{}{}
+
+	selfID := ps.srvr.Self().ID().String()
+
+	dkgIDs, err := ps.gov.DKGSet(round)
+	if err != nil {
+		log.Error("get dkg set fail", "round", round, "err", err)
+	}
+
+	// build dkg connection
+	_, inDKGSet := dkgIDs[selfID]
+	if inDKGSet {
+		delete(dkgIDs, selfID)
+		dkgLabel := peerLabel{set: dkgset, round: round}
+		for id := range dkgIDs {
+			ps.addDirectPeer(id, dkgLabel)
+		}
+	}
+
+	var inOneNotarySet bool
+	for cid := uint32(0); cid < ps.gov.GetNumChains(round); cid++ {
+		notaryIDs, err := ps.gov.NotarySet(round, cid)
+		if err != nil {
+			log.Error("get notary set fail",
+				"round", round, "chain id", cid, "err", err)
+			continue
+		}
+
+		label := peerLabel{set: notaryset, chainID: cid, round: round}
+		// not in notary set, add group
+		if _, ok := notaryIDs[selfID]; !ok {
+			var nodes []*enode.Node
+			for id := range notaryIDs {
+				nodes = append(nodes, ps.newNode(id))
+				ps.addLabel(id, label)
+			}
+			ps.srvr.AddGroup(notarySetName(cid, round), nodes, groupNodeNum)
+			continue
+		}
+
+		delete(notaryIDs, selfID)
+		for id := range notaryIDs {
+			ps.addDirectPeer(id, label)
+		}
+		inOneNotarySet = true
+	}
+
+	// build some connections to DKG nodes
+	if !inDKGSet && inOneNotarySet {
+		var nodes []*enode.Node
+		label := peerLabel{set: dkgset, round: round}
+		for id := range dkgIDs {
+			nodes = append(nodes, ps.newNode(id))
+			ps.addLabel(id, label)
+		}
+		ps.srvr.AddGroup(dkgSetName(round), nodes, groupNodeNum)
+	}
+}
+
+func (ps *peerSet) ForgetConnection(round uint64) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+	defer ps.dumpPeerLabel(fmt.Sprintf("ForgetConnection: %d", round))
+
+	for r := range ps.history {
+		if r <= round {
+			ps.forgetConnection(round)
+			delete(ps.history, r)
+		}
+	}
+}
+
+func (ps *peerSet) forgetConnection(round uint64) {
+	selfID := ps.srvr.Self().ID().String()
+	dkgIDs, err := ps.gov.DKGSet(round)
+	if err != nil {
+		log.Error("get dkg set fail", "round", round, "err", err)
+	}
+
+	_, inDKGSet := dkgIDs[selfID]
+	if inDKGSet {
+		delete(dkgIDs, selfID)
+		label := peerLabel{set: dkgset, round: round}
+		for id := range dkgIDs {
+			ps.removeDirectPeer(id, label)
+		}
+	}
+
+	var inOneNotarySet bool
+	for cid := uint32(0); cid < ps.gov.GetNumChains(round); cid++ {
+		notaryIDs, err := ps.gov.NotarySet(round, cid)
+		if err != nil {
+			log.Error("get notary set fail",
+				"round", round, "chain id", cid, "err", err)
+			continue
+		}
+
+		label := peerLabel{set: notaryset, chainID: cid, round: round}
+
+		// not in notary set, add group
+		if _, ok := notaryIDs[selfID]; !ok {
+			var nodes []*enode.Node
+			for id := range notaryIDs {
+				nodes = append(nodes, ps.newNode(id))
+				ps.removeLabel(id, label)
+			}
+			ps.srvr.RemoveGroup(notarySetName(cid, round))
+			continue
+		}
+
+		delete(notaryIDs, selfID)
+		for id := range notaryIDs {
+			ps.removeDirectPeer(id, label)
+		}
+		inOneNotarySet = true
+	}
+
+	// build some connections to DKG nodes
+	if !inDKGSet && inOneNotarySet {
+		var nodes []*enode.Node
+		label := peerLabel{set: dkgset, round: round}
+		for id := range dkgIDs {
+			nodes = append(nodes, ps.newNode(id))
+			ps.removeLabel(id, label)
+		}
+		ps.srvr.RemoveGroup(dkgSetName(round))
+	}
+}
+
 func (ps *peerSet) BuildNotaryConn(round uint64) {
 	ps.lock.Lock()
 	defer ps.lock.Unlock()
@@ -786,7 +926,7 @@ func (ps *peerSet) BuildNotaryConn(round uint64) {
 
 	ps.notaryHistory[round] = struct{}{}
 
-	selfID := ps.srvr.Self().ID.String()
+	selfID := ps.srvr.Self().ID().String()
 	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
 		s, err := ps.gov.NotarySet(round, chainID)
 		if err != nil {
@@ -797,7 +937,7 @@ func (ps *peerSet) BuildNotaryConn(round uint64) {
 
 		// not in notary set, add group
 		if _, ok := s[selfID]; !ok {
-			var nodes []*discover.Node
+			var nodes []*enode.Node
 			for id := range s {
 				nodes = append(nodes, ps.newNode(id))
 			}
@@ -843,7 +983,7 @@ func (ps *peerSet) ForgetNotaryConn(round uint64) {
 }
 
 func (ps *peerSet) forgetNotaryConn(round uint64) {
-	selfID := ps.srvr.Self().ID.String()
+	selfID := ps.srvr.Self().ID().String()
 	for chainID := uint32(0); chainID < ps.gov.GetNumChains(round); chainID++ {
 		s, err := ps.gov.NotarySet(round, chainID)
 		if err != nil {
@@ -872,11 +1012,15 @@ func notarySetName(chainID uint32, round uint64) string {
 	return fmt.Sprintf("%d-%d-notaryset", chainID, round)
 }
 
+func dkgSetName(round uint64) string {
+	return fmt.Sprintf("%d-dkgset", round)
+}
+
 func (ps *peerSet) BuildDKGConn(round uint64) {
 	ps.lock.Lock()
 	defer ps.lock.Unlock()
 	defer ps.dumpPeerLabel(fmt.Sprintf("BuildDKGConn: %d", round))
-	selfID := ps.srvr.Self().ID.String()
+	selfID := ps.srvr.Self().ID().String()
 	s, err := ps.gov.DKGSet(round)
 	if err != nil {
 		log.Error("get dkg set fail", "round", round)
@@ -912,7 +1056,7 @@ func (ps *peerSet) ForgetDKGConn(round uint64) {
 }
 
 func (ps *peerSet) forgetDKGConn(round uint64) {
-	selfID := ps.srvr.Self().ID.String()
+	selfID := ps.srvr.Self().ID().String()
 	s, err := ps.gov.DKGSet(round)
 	if err != nil {
 		log.Error("get dkg set fail", "round", round)
@@ -932,45 +1076,41 @@ func (ps *peerSet) forgetDKGConn(round uint64) {
 	}
 }
 
-// make sure the ps.lock is hold
+// make sure the ps.lock is held
 func (ps *peerSet) addDirectPeer(id string, label peerLabel) {
-	log.Trace("peerSet addDirectPeer", "id", id[:8], "round", label.round, "cid", label.chainID)
-	// if the peer exists add the label
-	if p, ok := ps.peers[id]; ok {
-		p.addLabel(label)
+	ps.addLabel(id, label)
+	ps.srvr.AddDirectPeer(ps.newNode(id))
+}
+
+// make sure the ps.lock is held
+func (ps *peerSet) removeDirectPeer(id string, label peerLabel) {
+	ps.removeLabel(id, label)
+	if len(ps.peer2Labels[id]) == 0 {
+		ps.srvr.RemoveDirectPeer(ps.newNode(id))
 	}
+}
 
+// make sure the ps.lock is held
+func (ps *peerSet) addLabel(id string, label peerLabel) {
 	if _, ok := ps.peer2Labels[id]; !ok {
 		ps.peer2Labels[id] = make(map[peerLabel]struct{})
 	}
-
 	if _, ok := ps.label2Peers[label]; !ok {
 		ps.label2Peers[label] = make(map[string]struct{})
 	}
-
 	ps.peer2Labels[id][label] = struct{}{}
 	ps.label2Peers[label][id] = struct{}{}
-	ps.srvr.AddDirectPeer(ps.newNode(id))
 }
 
-// make sure the ps.lock is hold
-func (ps *peerSet) removeDirectPeer(id string, label peerLabel) {
-	if p, ok := ps.peers[id]; ok {
-		p.removeLabel(label)
-	}
-
+// make sure the ps.lock is held
+func (ps *peerSet) removeLabel(id string, label peerLabel) {
 	delete(ps.peer2Labels[id], label)
-
+	delete(ps.label2Peers[label], id)
 	if len(ps.peer2Labels[id]) == 0 {
-		ps.srvr.RemoveDirectPeer(ps.newNode(id))
 		delete(ps.peer2Labels, id)
 	}
-
-	if _, ok := ps.label2Peers[label]; ok {
-		delete(ps.label2Peers[label], id)
-		if len(ps.label2Peers[label]) == 0 {
-			delete(ps.label2Peers, label)
-		}
+	if len(ps.label2Peers[label]) == 0 {
+		delete(ps.label2Peers, label)
 	}
 }