[PATCH 2/2] the connection should less than max connet peers

If the connection more than max connect peers, disconnect one node. Signed-off-by: Jin Qing <[email protected]>
DNAProject · Sep 19, 2017 · eb1b9ef · eb1b9ef
1 parent 88c6fd9
commit eb1b9ef
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/net/node/infoUpdate.go b/net/node/infoUpdate.go
@@ -4,6 +4,7 @@ import (
 	"DNA/common/config"
 	"DNA/common/log"
 	"DNA/core/ledger"
+	"DNA/events"
 	. "DNA/net/message"
 	. "DNA/net/protocol"
 	"math/rand"
@@ -243,6 +244,14 @@ func (node *node) updateNodeInfo() {
 	//close(quit)
 }
 
+func (node *node) CheckConnCnt() {
+	//compare if connect count is larger than DefaultMaxPeers, disconnect one of the connection
+	if node.nbrNodes.GetConnectionCnt() > node.GetDefaultMaxPeers() {
+		disconnNode := node.RandGetANbr()
+		node.eventQueue.GetEvent("disconnect").Notify(events.EventNodeDisconnect, disconnNode)
+	}
+}
+
 func (node *node) updateConnection() {
 	t := time.NewTimer(time.Second * CONNMONITOR)
 	for {
@@ -251,6 +260,7 @@ func (node *node) updateConnection() {
 			node.ConnectSeeds()
 			//node.TryConnect()
 			node.ConnectNode()
+			node.CheckConnCnt()
 			t.Stop()
 			t.Reset(time.Second * CONNMONITOR)
 		}

diff --git a/net/node/nodeMap.go b/net/node/nodeMap.go
@@ -154,3 +154,12 @@ func (node *node) GetNbrNodeCnt() uint32 {
 	}
 	return count
 }
+
+func (node *node) RandGetANbr() Noder {
+	for _, n := range node.nbrNodes.List {
+		if n.GetState() == ESTABLISH {
+			return n
+		}
+	}
+	return nil
+}