Rework PeerListSupport a little bit

Author: Petra Bierleutgeb

src/main/scala/io/iohk/ethereum/blockchain/sync/Blacklist.scala

@@ -13,19 +13,16 @@ import scala.jdk.DurationConverters._
 import Blacklist.BlacklistId
 
 trait Blacklist {
-
   def isBlacklisted(id: BlacklistId): Boolean
   def add(id: BlacklistId, duration: FiniteDuration, reason: String): Unit
   def remove(id: BlacklistId): Unit
   def keys: Set[BlacklistId]
 }
 
 object Blacklist {
-
   trait BlacklistId {
     def value: String
   }
-
 }
 
 final case class CacheBasedBlacklist(cache: Cache[BlacklistId, String])

src/main/scala/io/iohk/ethereum/blockchain/sync/PeerListSupport.scala

@@ -11,6 +11,7 @@ import io.iohk.ethereum.utils.Config.SyncConfig
 import scala.concurrent.duration._
 import scala.concurrent.ExecutionContext.Implicits.global
 
+// will be removed once regular sync is switched to new blacklist/peerlist implementation
 trait PeerListSupport {
   self: Actor with ActorLogging with BlacklistSupport =>
   import PeerListSupport._

src/main/scala/io/iohk/ethereum/blockchain/sync/PeerListSupportNg.scala

@@ -8,56 +8,64 @@ import io.iohk.ethereum.network.PeerEventBusActor.{PeerSelector, Subscribe, Unsu
 import io.iohk.ethereum.network.{EtcPeerManagerActor, Peer, PeerId}
 import io.iohk.ethereum.utils.Config.SyncConfig
 
+import scala.concurrent.ExecutionContext
 import scala.concurrent.duration._
-import scala.concurrent.ExecutionContext.Implicits.global // FIXME copied from current impl but not sure global EC is a good choice
 
-trait PeerListSupportNg {
-  self: Actor with ActorLogging =>
+trait PeerListSupportNg { self: Actor with ActorLogging =>
   import PeerListSupportNg._
 
+  private implicit val ec: ExecutionContext = context.dispatcher
+
   def etcPeerManager: ActorRef
   def peerEventBus: ActorRef
   def blacklist: Blacklist
   def syncConfig: SyncConfig
   def scheduler: Scheduler
 
-  var handshakedPeers: PeersMap = Map.empty
+  protected var handshakedPeers: Map[PeerId, PeerWithInfo] = Map.empty
 
   scheduler.scheduleWithFixedDelay(
     0.seconds,
     syncConfig.peersScanInterval,
     etcPeerManager,
     EtcPeerManagerActor.GetHandshakedPeers
-  )(global, context.self)
+  )
 
-  def removePeer(peerId: PeerId): Unit = {
-    peerEventBus ! Unsubscribe(PeerDisconnectedClassifier(PeerSelector.WithId(peerId)))
-    handshakedPeers.find(_._1.id == peerId).foreach { case (peer, _) => blacklist.remove(peer.id) }
-    handshakedPeers = handshakedPeers.filterNot(_._1.id == peerId)
+  def handlePeerListMessages: Receive = {
+    case EtcPeerManagerActor.HandshakedPeers(peers) => updatePeers(peers)
+    case PeerDisconnected(peerId) => removePeer(peerId)
   }
 
-  def peersToDownloadFrom: PeersMap =
-    handshakedPeers.filterNot { case (p, _) => blacklist.isBlacklisted(p.id) }
+  def peersToDownloadFrom: Map[PeerId, PeerWithInfo] =
+    handshakedPeers.filterNot { case (peerId, _) =>
+      blacklist.isBlacklisted(peerId)
+    }
 
-  def handlePeerListMessages: Receive = {
-    case EtcPeerManagerActor.HandshakedPeers(peers) =>
-      peers.keys.filterNot(handshakedPeers.contains).foreach { peer =>
-        peerEventBus ! Subscribe(PeerDisconnectedClassifier(PeerSelector.WithId(peer.id)))
-      }
-      handshakedPeers = peers
-
-    case PeerDisconnected(peerId) if handshakedPeers.exists(_._1.id == peerId) =>
-      removePeer(peerId)
+  def peerById(peerId: PeerId): Option[Peer] = handshakedPeers.get(peerId).map(_.peer)
+
+  def blacklistIfHandshaked(peerId: PeerId, duration: FiniteDuration, reason: String): Unit =
+    handshakedPeers.get(peerId).foreach(_ => blacklist.add(peerId, duration, reason))
+
+  private def updatePeers(peers: Map[Peer, PeerInfo]): Unit = {
+    val updated = peers.map { case (peer, peerInfo) =>
+      (peer.id, PeerWithInfo(peer, peerInfo))
+    }
+    updated.filterNot(p => handshakedPeers.keySet.contains(p._1)).foreach { case (peerId, _) =>
+      peerEventBus ! Subscribe(PeerDisconnectedClassifier(PeerSelector.WithId(peerId)))
+    }
+    handshakedPeers = updated
   }
 
-  def peerById(peerId: PeerId): Option[Peer] = handshakedPeers collectFirst {
-    case (peer, _) if peer.id == peerId => peer
+  private def removePeer(peerId: PeerId): Unit = {
+    if (handshakedPeers.keySet.contains(peerId)) {
+      peerEventBus ! Unsubscribe(PeerDisconnectedClassifier(PeerSelector.WithId(peerId)))
+      blacklist.remove(peerId)
+      handshakedPeers = handshakedPeers - peerId
+    }
   }
 
-  def blacklistIfHandshaked(peer: Peer, duration: FiniteDuration, reason: String): Unit =
-    handshakedPeers.get(peer).foreach(_ => blacklist.add(peer.id, duration, reason))
 }
 
 object PeerListSupportNg {
-  type PeersMap = Map[Peer, PeerInfo]
+  final case class PeerWithInfo(peer: Peer, peerInfo: PeerInfo)
 }

src/main/scala/io/iohk/ethereum/blockchain/sync/fast/FastSync.scala

@@ -4,6 +4,7 @@ import java.time.Instant
 import akka.actor._
 import akka.util.ByteString
 import cats.data.NonEmptyList
+import io.iohk.ethereum.blockchain.sync.PeerListSupportNg.PeerWithInfo
 import io.iohk.ethereum.blockchain.sync.PeerRequestHandler.ResponseReceived
 import io.iohk.ethereum.blockchain.sync.SyncProtocol.Status.Progress
 import io.iohk.ethereum.blockchain.sync._
@@ -20,7 +21,7 @@ import io.iohk.ethereum.db.storage.{AppStateStorage, FastSyncStateStorage}
 import io.iohk.ethereum.domain._
 import io.iohk.ethereum.mpt.MerklePatriciaTrie
 import io.iohk.ethereum.network.EtcPeerManagerActor.PeerInfo
-import io.iohk.ethereum.network.Peer
+import io.iohk.ethereum.network.{Peer, PeerId}
 import io.iohk.ethereum.network.p2p.messages.Codes
 import io.iohk.ethereum.network.p2p.messages.PV62._
 import io.iohk.ethereum.network.p2p.messages.PV63._
@@ -541,7 +542,7 @@ class FastSync(
       requestedReceipts = requestedReceipts - handler
 
       requestedHeaders -= peer
-      if (handshakedPeers.contains(peer)) {
+      if (handshakedPeers.contains(peer.id)) {
         blacklist.add(peer.id, blacklistDuration, reason)
       }
     }
@@ -567,7 +568,8 @@ class FastSync(
     }
 
     private def printStatus(): Unit = {
-      val formatPeer: Peer => String = peer =>
+      def formatPeerEntry(entry: PeerWithInfo): String = formatPeer(entry.peer)
+      def formatPeer(peer: Peer): String =
         s"${peer.remoteAddress.getAddress.getHostAddress}:${peer.remoteAddress.getPort}"
       val blacklistedIds = blacklist.keys
       log.info(
@@ -583,7 +585,7 @@ class FastSync(
             | blacklisted({})
             |""".stripMargin.replace("\n", " "),
         assignedHandlers.values.map(formatPeer).toSeq.sorted.mkString(", "),
-        handshakedPeers.keys.map(formatPeer).toSeq.sorted.mkString(", "),
+        handshakedPeers.values.toList.map(e => formatPeerEntry(e)).sorted.mkString(", "),
         blacklistedIds.map(_.value).mkString(", ")
       )
     }
@@ -609,12 +611,12 @@ class FastSync(
     }
 
     private def getPeersWithFreshEnoughPivot(
-        peers: NonEmptyList[(Peer, PeerInfo)],
+        peers: NonEmptyList[PeerWithInfo],
         state: SyncState,
         syncConfig: SyncConfig
     ): List[(Peer, BigInt)] = {
       peers.collect {
-        case (peer, info) if hasBestBlockFreshEnoughToUpdatePivotBlock(info, state, syncConfig) =>
+        case PeerWithInfo(peer, info) if hasBestBlockFreshEnoughToUpdatePivotBlock(info, state, syncConfig) =>
           (peer, info.maxBlockNumber)
       }
     }
@@ -626,17 +628,17 @@ class FastSync(
       !(syncState.updatingPivotBlock || stateSyncRestartRequested)
 
     def pivotBlockIsStale(): Boolean = {
-      val currentPeers = peersToDownloadFrom.toList
-      if (currentPeers.isEmpty) {
+      val peersWithInfo = peersToDownloadFrom.values.toList
+      if (peersWithInfo.isEmpty) {
         false
       } else {
-        val peerWithBestBlockInNetwork = currentPeers.maxBy(peerWithNum => peerWithNum._2.maxBlockNumber)
+        val peerWithBestBlockInNetwork = peersWithInfo.maxBy(_.peerInfo.maxBlockNumber)
 
         val bestPossibleTargetDifferenceInNetwork =
-          (peerWithBestBlockInNetwork._2.maxBlockNumber - syncConfig.pivotBlockOffset) - syncState.pivotBlock.number
+          (peerWithBestBlockInNetwork.peerInfo.maxBlockNumber - syncConfig.pivotBlockOffset) - syncState.pivotBlock.number
 
         val peersWithTooFreshPossiblePivotBlock =
-          getPeersWithFreshEnoughPivot(NonEmptyList.fromListUnsafe(currentPeers), syncState, syncConfig)
+          getPeersWithFreshEnoughPivot(NonEmptyList.fromListUnsafe(peersWithInfo), syncState, syncConfig)
 
         if (peersWithTooFreshPossiblePivotBlock.isEmpty) {
           log.info(
@@ -652,7 +654,7 @@ class FastSync(
             "There are {} peers with possible new pivot block, " +
               "best known pivot in current peer list has number {}",
             peersWithTooFreshPossiblePivotBlock.size,
-            peerWithBestBlockInNetwork._2.maxBlockNumber
+            peerWithBestBlockInNetwork.peerInfo.maxBlockNumber
           )
 
           pivotBlockIsStale
@@ -715,7 +717,7 @@ class FastSync(
           .filter(p => peerRequestsTime.get(p.peer).forall(d => d.plusMillis(fastSyncThrottle.toMillis).isBefore(now)))
         peers
           .take(maxConcurrentRequests - assignedHandlers.size)
-          .sortBy(_.info.maxBlockNumber)(Ordering[BigInt].reverse)
+          .sortBy(_.peerInfo.maxBlockNumber)(Ordering[BigInt].reverse)
           .foreach(assignBlockchainWork)
       }
     }
@@ -803,8 +805,10 @@ class FastSync(
       peerRequestsTime += (peer -> Instant.now())
     }
 
-    def unassignedPeers: List[PeerWithInfo] =
-      (peersToDownloadFrom -- assignedHandlers.values).map(PeerWithInfo.tupled).toList
+    def unassignedPeers: List[PeerWithInfo] = {
+      val assignedPeers = assignedHandlers.values.map(_.id).toList
+      peersToDownloadFrom.removedAll(assignedPeers).values.toList
+    }
 
     def blockchainDataToDownload: Boolean =
       syncState.blockChainWorkQueued || syncState.bestBlockHeaderNumber < syncState.safeDownloadTarget
@@ -838,8 +842,6 @@ class FastSync(
 
 object FastSync {
 
-  case class PeerWithInfo(peer: Peer, info: PeerInfo)
-
   // scalastyle:off parameter.number
   def props(
       fastSyncStateStorage: FastSyncStateStorage,