Try to choose candidates from different groups

hyperledger-labs · Dec 24, 2020 · d7efdbf · d7efdbf
1 parent 5d8be6d
commit d7efdbf
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/main/scala/scorex/core/network/peer/PeerManager.scala b/src/main/scala/scorex/core/network/peer/PeerManager.scala
@@ -149,15 +149,24 @@ object PeerManager {
                           sc: ScorexContext): Map[InetSocketAddress, PeerInfo] = knownPeers
     }
 
+    private  def getIpGroup(addr : InetSocketAddress) : Int = {
+      val ip = addr.getAddress.getAddress
+      val group = ((ip(0) & 0xFF) << 8) | (ip(1) & 0xFF)
+      group
+    }
+
     case class RandomPeerExcluding(excludedPeers: Seq[PeerInfo]) extends GetPeers[Option[PeerInfo]] {
 
       override def choose(knownPeers: Map[InetSocketAddress, PeerInfo],
                           blacklistedPeers: Seq[InetAddress],
                           sc: ScorexContext): Option[PeerInfo] = {
-        val candidates = knownPeers.values.filterNot { p =>
-          excludedPeers.exists(_.peerSpec.address == p.peerSpec.address) &&
-            blacklistedPeers.exists(addr => p.peerSpec.address.map(_.getAddress).contains(addr))
+        val excludedGroups = excludedPeers.flatMap(_.peerSpec.address).map(getIpGroup(_)).toSet
+        val allCandidates = knownPeers.values.filterNot { p =>
+          excludedPeers.exists(_.peerSpec.address == p.peerSpec.address) ||
+          blacklistedPeers.exists(addr => p.peerSpec.address.map(_.getAddress).contains(addr))
         }.toSeq
+        val preferredCandidates = allCandidates.filterNot(_.peerSpec.address.fold(true)(addr => excludedGroups.contains(getIpGroup(addr))))
+        val candidates = if (preferredCandidates.nonEmpty) preferredCandidates else allCandidates
         if (candidates.nonEmpty) Some(candidates(Random.nextInt(candidates.size)))
         else None
       }