EC-544 Improve pending transaction manager performance

Michal Mrozek · Michal Mrozek · commit 3d6808f121e5 · 2018-07-10T10:23:13.000+02:00
diff --git a/src/main/scala/io/iohk/ethereum/blockchain/sync/RegularSync.scala b/src/main/scala/io/iohk/ethereum/blockchain/sync/RegularSync.scala
@@ -361,7 +361,7 @@ class RegularSync(
   private def processBlockHeaders(peer: Peer, headers: Seq[BlockHeader]) = ledger.resolveBranch(headers) match {
     case NewBetterBranch(oldBranch) =>
       val transactionsToAdd = oldBranch.flatMap(_.body.transactionList)
-      pendingTransactionsManager ! PendingTransactionsManager.AddTransactions(transactionsToAdd.toList)
+      pendingTransactionsManager ! PendingTransactionsManager.AddTransactions(transactionsToAdd.toSet)
       val hashes = headers.take(blockBodiesPerRequest).map(_.hash)
       requestBlockBodies(peer, GetBlockBodies(hashes))
 
@@ -505,7 +505,7 @@ class RegularSync(
 
   private def updateTxAndOmmerPools(blocksAdded: Seq[Block], blocksRemoved: Seq[Block]): Unit = {
     blocksRemoved.headOption.foreach(block => ommersPool ! AddOmmers(block.header))
-    blocksRemoved.foreach(block => pendingTransactionsManager ! AddTransactions(block.body.transactionList.toList))
+    blocksRemoved.foreach(block => pendingTransactionsManager ! AddTransactions(block.body.transactionList.toSet))
 
     blocksAdded.foreach { block =>
       ommersPool ! RemoveOmmers(block.header :: block.body.uncleNodesList.toList)
diff --git a/src/main/scala/io/iohk/ethereum/transactions/PendingTransactionsManager.scala b/src/main/scala/io/iohk/ethereum/transactions/PendingTransactionsManager.scala
@@ -2,26 +2,27 @@ package io.iohk.ethereum.transactions
 
 import akka.actor.{Actor, ActorRef, Cancellable, Props}
 import akka.util.{ByteString, Timeout}
+import com.google.common.cache.{Cache, CacheBuilder, RemovalCause, RemovalNotification}
 import io.iohk.ethereum.domain.SignedTransaction
 import io.iohk.ethereum.network.PeerEventBusActor.PeerEvent.MessageFromPeer
 import io.iohk.ethereum.network.PeerEventBusActor.SubscriptionClassifier.MessageClassifier
 import io.iohk.ethereum.network.PeerEventBusActor.{PeerEvent, PeerSelector, Subscribe, SubscriptionClassifier}
 import io.iohk.ethereum.network.PeerManagerActor.Peers
-import io.iohk.ethereum.network.{EtcPeerManagerActor, Peer, PeerId, PeerManagerActor}
 import io.iohk.ethereum.network.p2p.messages.CommonMessages.SignedTransactions
+import io.iohk.ethereum.network.{EtcPeerManagerActor, Peer, PeerId, PeerManagerActor}
 import io.iohk.ethereum.utils.TxPoolConfig
-
-import scala.concurrent.duration._
+import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration._
 
 object PendingTransactionsManager {
   def props(txPoolConfig: TxPoolConfig, peerManager: ActorRef, etcPeerManager: ActorRef, peerMessageBus: ActorRef): Props =
     Props(new PendingTransactionsManager(txPoolConfig, peerManager, etcPeerManager, peerMessageBus))
 
-  case class AddTransactions(signedTransactions: List[SignedTransaction])
+  case class AddTransactions(signedTransactions: Set[SignedTransaction])
 
   object AddTransactions{
-    def apply(txs: SignedTransaction*): AddTransactions = AddTransactions(txs.toList)
+    def apply(txs: SignedTransaction*): AddTransactions = AddTransactions(txs.toSet)
   }
 
   case class AddOrOverrideTransaction(signedTransaction: SignedTransaction)
@@ -39,21 +40,28 @@ object PendingTransactionsManager {
 }
 
 class PendingTransactionsManager(txPoolConfig: TxPoolConfig, peerManager: ActorRef,
-                                 etcPeerManager: ActorRef, peerEventBus: ActorRef) extends Actor {
+                                  etcPeerManager: ActorRef, peerEventBus: ActorRef) extends Actor {
 
   import PendingTransactionsManager._
   import akka.pattern.ask
 
-  /**
-    * stores all pending transactions
-   */
-  var pendingTransactions: List[PendingTransaction] = Nil
-
   /**
     * stores information which tx hashes are "known" by which peers
     */
   var knownTransactions: Map[ByteString, Set[PeerId]] = Map.empty
 
+  /**
+    * stores all pending transactions
+    */
+  val pendingTransactions: Cache[ByteString, PendingTransaction] = CacheBuilder.newBuilder()
+    .expireAfterWrite(txPoolConfig.transactionTimeout._1, txPoolConfig.transactionTimeout._2)
+    .maximumSize(txPoolConfig.txPoolSize)
+    .removalListener(
+      (notification: RemovalNotification[ByteString, PendingTransaction]) => if (notification.getCause == RemovalCause.EXPIRED) {
+        knownTransactions = knownTransactions.filterNot(_._1 == notification.getKey)
+      }
+    ).build()
+
   /**
     * stores transactions timeouts by tx hash
     */
@@ -67,68 +75,61 @@ class PendingTransactionsManager(txPoolConfig: TxPoolConfig, peerManager: ActorR
   // scalastyle:off method.length
   override def receive: Receive = {
     case PeerEvent.PeerHandshakeSuccessful(peer, _) =>
-      self ! NotifyPeer(pendingTransactions.map(_.stx), peer)
+      pendingTransactions.cleanUp()
+      val stxs = pendingTransactions.asMap().values().asScala.toSeq.map(_.stx)
+      self ! NotifyPeer(stxs, peer)
 
     case AddTransactions(signedTransactions) =>
-      val transactionsToAdd = signedTransactions.filterNot(t => pendingTransactions.map(_.stx).contains(t))
+      pendingTransactions.cleanUp()
+      val stxs = pendingTransactions.asMap().values().asScala.map(_.stx).toSet
+      val transactionsToAdd = signedTransactions.diff(stxs)
       if (transactionsToAdd.nonEmpty) {
-        transactionsToAdd.foreach(setTimeout)
         val timestamp = System.currentTimeMillis()
-        pendingTransactions = (transactionsToAdd.map(PendingTransaction(_, timestamp)) ++ pendingTransactions).take(txPoolConfig.txPoolSize)
+        transactionsToAdd.foreach(t => pendingTransactions.put(t.hash, PendingTransaction(t, timestamp)))
         (peerManager ? PeerManagerActor.GetPeers).mapTo[Peers].foreach { peers =>
-          peers.handshaked.foreach { peer => self ! NotifyPeer(transactionsToAdd, peer) }
+          peers.handshaked.foreach { peer => self ! NotifyPeer(transactionsToAdd.toSeq, peer) }
         }
       }
 
     case AddOrOverrideTransaction(newStx) =>
-      val (obsoleteTxs, txsWithoutObsoletes) = pendingTransactions.partition(ptx =>
-        ptx.stx.senderAddress == newStx.senderAddress &&
-        ptx.stx.tx.nonce == newStx.tx.nonce)
-      obsoleteTxs.map(_.stx).foreach(clearTimeout)
+      pendingTransactions.cleanUp()
+      val obsoleteTxs = pendingTransactions.asMap().asScala.filter(
+        ptx => ptx._2.stx.senderAddress == newStx.senderAddress && ptx._2.stx.tx.nonce == newStx.tx.nonce
+      )
+      pendingTransactions.invalidateAll(obsoleteTxs.keys.asJava)
 
       val timestamp = System.currentTimeMillis()
-      pendingTransactions = (PendingTransaction(newStx, timestamp) +: txsWithoutObsoletes).take(txPoolConfig.txPoolSize)
-      setTimeout(newStx)
+      pendingTransactions.put(newStx.hash, PendingTransaction(newStx, timestamp))
 
-      (peerManager ? PeerManagerActor.GetPeers).mapTo[Peers].foreach { peers =>
-        peers.handshaked.foreach { peer => self ! NotifyPeer(List(newStx), peer) }
+      (peerManager ? PeerManagerActor.GetPeers).mapTo[Peers].foreach {
+        peers => peers.handshaked.foreach { peer => self ! NotifyPeer(List(newStx), peer) }
       }
 
     case NotifyPeer(signedTransactions, peer) =>
+      pendingTransactions.cleanUp()
       val txsToNotify = signedTransactions
-        .filter(stx => pendingTransactions.exists(_.stx.hash == stx.hash)) // signed transactions that are still pending
+        .filter(stx => pendingTransactions.asMap().containsKey(stx.hash)) // signed transactions that are still pending
         .filterNot(isTxKnown(_, peer.id)) // and not known by peer
 
-        if (txsToNotify.nonEmpty) {
-          etcPeerManager ! EtcPeerManagerActor.SendMessage(SignedTransactions(txsToNotify), peer.id)
-          txsToNotify.foreach(setTxKnown(_, peer.id))
-        }
+      if (txsToNotify.nonEmpty) {
+        etcPeerManager ! EtcPeerManagerActor.SendMessage(SignedTransactions(txsToNotify), peer.id)
+        txsToNotify.foreach(setTxKnown(_, peer.id))
+      }
 
     case GetPendingTransactions =>
-      sender() ! PendingTransactionsResponse(pendingTransactions)
+      pendingTransactions.cleanUp()
+      sender() ! PendingTransactionsResponse(pendingTransactions.asMap().asScala.values.toSeq)
 
     case RemoveTransactions(signedTransactions) =>
-      pendingTransactions = pendingTransactions.filterNot(pt => signedTransactions.contains(pt.stx))
-      knownTransactions = knownTransactions.filterNot(signedTransactions.map(_.hash).contains)
-      signedTransactions.foreach(clearTimeout)
+      pendingTransactions.invalidateAll(signedTransactions.map(_.hash).asJava)
+      knownTransactions = knownTransactions -- signedTransactions.map(_.hash)
 
     case MessageFromPeer(SignedTransactions(signedTransactions), peerId) =>
-      self ! AddTransactions(signedTransactions.toList)
+      self ! AddTransactions(signedTransactions.toSet)
       signedTransactions.foreach(setTxKnown(_, peerId))
 
     case ClearPendingTransactions =>
-      pendingTransactions = Nil
-  }
-
-  private def setTimeout(stx: SignedTransaction): Unit = {
-    timeouts.get(stx.hash).map(_.cancel())
-    val cancellable = context.system.scheduler.scheduleOnce(txPoolConfig.transactionTimeout, self, RemoveTransactions(Seq(stx)))
-    timeouts += (stx.hash -> cancellable)
-  }
-
-  private def clearTimeout(stx: SignedTransaction): Unit = {
-    timeouts.get(stx.hash).map(_.cancel())
-    timeouts -= stx.hash
+      pendingTransactions.invalidateAll()
   }
 
   private def isTxKnown(signedTransaction: SignedTransaction, peerId: PeerId): Boolean =
diff --git a/src/test/scala/io/iohk/ethereum/blockchain/sync/RegularSyncSpec.scala b/src/test/scala/io/iohk/ethereum/blockchain/sync/RegularSyncSpec.scala
@@ -73,7 +73,7 @@ class RegularSyncSpec extends TestKit(ActorSystem("RegularSync_system")) with Wo
         sendNewBlockMsg(newBlock)
 
         ommersPool.expectMsg(AddOmmers(List(oldBlock.header)))
-        txPool.expectMsg(AddTransactions(oldBlock.body.transactionList.toList))
+        txPool.expectMsg(AddTransactions(oldBlock.body.transactionList.toSet))
 
         ommersPool.expectMsg(RemoveOmmers(newBlock.header :: newBlock.body.uncleNodesList.toList))
         txPool.expectMsg(RemoveTransactions(newBlock.body.transactionList.toList))
@@ -213,7 +213,7 @@ class RegularSyncSpec extends TestKit(ActorSystem("RegularSync_system")) with Wo
         sendMinedBlockMsg(newBlock)
 
         ommersPool.expectMsg(AddOmmers(List(oldBlock.header)))
-        txPool.expectMsg(AddTransactions(oldBlock.body.transactionList.toList))
+        txPool.expectMsg(AddTransactions(oldBlock.body.transactionList.toSet))
 
         ommersPool.expectMsg(RemoveOmmers(newBlock.header :: newBlock.body.uncleNodesList.toList))
         txPool.expectMsg(RemoveTransactions(newBlock.body.transactionList.toList))
@@ -268,7 +268,7 @@ class RegularSyncSpec extends TestKit(ActorSystem("RegularSync_system")) with Wo
 
         etcPeerManager.expectMsg(EtcPeerManagerActor.GetHandshakedPeers)
         etcPeerManager.expectMsg(EtcPeerManagerActor.SendMessage(GetBlockBodies(newBlocks.map(_.header.hash)), peer1.id))
-        txPool.expectMsg(AddTransactions(oldBlocks.flatMap(_.body.transactionList).toList))
+        txPool.expectMsg(AddTransactions(oldBlocks.flatMap(_.body.transactionList).toSet))
         ommersPool.expectMsg(AddOmmers(oldBlocks.head.header))
       }
 
diff --git a/src/test/scala/io/iohk/ethereum/transactions/PendingTransactionsManagerSpec.scala b/src/test/scala/io/iohk/ethereum/transactions/PendingTransactionsManagerSpec.scala
@@ -1,7 +1,6 @@
 package io.iohk.ethereum.transactions
 
 import java.net.InetSocketAddress
-
 import akka.actor.ActorSystem
 import akka.pattern.ask
 import akka.testkit.TestProbe
@@ -21,7 +20,6 @@ import io.iohk.ethereum.utils.TxPoolConfig
 import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.{FlatSpec, Matchers}
 import org.bouncycastle.crypto.AsymmetricCipherKeyPair
-
 import scala.concurrent.duration._
 
 class PendingTransactionsManagerSpec extends FlatSpec with Matchers with ScalaFutures with NormalPatience {
@@ -70,28 +68,33 @@ class PendingTransactionsManagerSpec extends FlatSpec with Matchers with ScalaFu
     pendingTransactionsManager ! MessageFromPeer(msg1, peer1.id)
     peerManager.expectMsg(PeerManagerActor.GetPeers)
     peerManager.reply(Peers(Map(peer1 -> Handshaked, peer2 -> Handshaked, peer3 -> Handshaked)))
-    etcPeerManager.expectMsgAllOf(
-      EtcPeerManagerActor.SendMessage(SignedTransactions(msg1.txs), peer2.id),
-      EtcPeerManagerActor.SendMessage(SignedTransactions(msg1.txs), peer3.id)
+
+    val resps1 = etcPeerManager.expectMsgAllConformingOf(
+      classOf[EtcPeerManagerActor.SendMessage], classOf[EtcPeerManagerActor.SendMessage]
     )
+
+    resps1.map(_.peerId) should contain allOf (peer2.id, peer3.id)
+    resps1.map(_.message.underlyingMsg).foreach { case SignedTransactions(txs) => txs.toSet shouldEqual msg1.txs.toSet }
     etcPeerManager.expectNoMsg()
 
     val msg2 = SignedTransactions(Seq.fill(5)(newStx()))
     pendingTransactionsManager ! MessageFromPeer(msg2, peer2.id)
     peerManager.expectMsg(PeerManagerActor.GetPeers)
     peerManager.reply(Peers(Map(peer1 -> Handshaked, peer2 -> Handshaked, peer3 -> Handshaked)))
-    etcPeerManager.expectMsgAllOf(
-      EtcPeerManagerActor.SendMessage(SignedTransactions(msg2.txs), peer1.id),
-      EtcPeerManagerActor.SendMessage(SignedTransactions(msg2.txs), peer3.id)
+
+    val resps2 = etcPeerManager.expectMsgAllConformingOf(
+      classOf[EtcPeerManagerActor.SendMessage], classOf[EtcPeerManagerActor.SendMessage]
     )
+    resps2.map(_.peerId) should contain allOf (peer1.id, peer3.id)
+    resps2.map(_.message.underlyingMsg).foreach { case SignedTransactions(txs) => txs.toSet shouldEqual msg2.txs.toSet }
     etcPeerManager.expectNoMsg()
 
     pendingTransactionsManager ! RemoveTransactions(msg1.txs.dropRight(4))
     pendingTransactionsManager ! RemoveTransactions(msg2.txs.drop(2))
 
     val pendingTxs = (pendingTransactionsManager ? GetPendingTransactions).mapTo[PendingTransactionsResponse].futureValue
     pendingTxs.pendingTransactions.size shouldBe 6
-    pendingTxs.pendingTransactions.map(_.stx) shouldBe msg2.txs.take(2) ++ msg1.txs.takeRight(4)
+    pendingTxs.pendingTransactions.map(_.stx).toSet shouldBe (msg2.txs.take(2) ++ msg1.txs.takeRight(4)).toSet
   }
 
   it should "not add pending transaction again when it was removed while waiting for peers" in new TestSetup {
@@ -132,7 +135,7 @@ class PendingTransactionsManagerSpec extends FlatSpec with Matchers with ScalaFu
     val pendingTxs = (pendingTransactionsManager ? GetPendingTransactions).mapTo[PendingTransactionsResponse]
       .futureValue.pendingTransactions
 
-    pendingTxs.map(_.stx) shouldEqual List(overrideTx, otherTx)
+    pendingTxs.map(_.stx).toSet shouldEqual Set(overrideTx, otherTx)
 
     // overriden TX will still be broadcast to peers
     etcPeerManager.expectMsgAllOf(
