add quarantinedButHarmless check for tests

apache · Nov 8, 2024 · be0ccc7 · be0ccc7
1 parent 38dbe25
commit be0ccc7
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/remote/src/main/scala/org/apache/pekko/remote/artery/ArteryTransport.scala b/remote/src/main/scala/org/apache/pekko/remote/artery/ArteryTransport.scala
@@ -108,9 +108,9 @@ private[remote] object AssociationState {
       quarantined = ImmutableLongMap.empty[QuarantinedTimestamp],
       new AtomicReference(UniqueRemoteAddressValue(None, Nil)))
 
-  final case class QuarantinedTimestamp(nanoTime: Long) {
+  final case class QuarantinedTimestamp(nanoTime: Long, harmless: Boolean = false) {
     override def toString: String =
-      s"Quarantined ${TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - nanoTime)} seconds ago"
+      s"Quarantined ${TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - nanoTime)} seconds ago (harmless=$harmless)"
   }
 
   private final case class UniqueRemoteAddressValue(
@@ -159,6 +159,13 @@ private[remote] final class AssociationState private (
 
   def isQuarantined(uid: Long): Boolean = quarantined.contains(uid)
 
+  def quarantinedButHarmless(uid: Long): Boolean = {
+    quarantined.get(uid) match {
+      case OptionVal.Some(qt) => qt.harmless
+      case _                  => false
+    }
+  }
+
   @tailrec def completeUniqueRemoteAddress(peer: UniqueAddress): Unit = {
     val current = _uniqueRemoteAddress.get()
     if (current.uniqueRemoteAddress.isEmpty) {
@@ -196,14 +203,14 @@ private[remote] final class AssociationState private (
       quarantined,
       new AtomicReference(UniqueRemoteAddressValue(Some(remoteAddress), Nil)))
 
-  def newQuarantined(): AssociationState =
+  def newQuarantined(harmless: Boolean = false): AssociationState =
     uniqueRemoteAddress() match {
       case Some(a) =>
         new AssociationState(
           incarnation,
           lastUsedTimestamp = new AtomicLong(System.nanoTime()),
           controlIdleKillSwitch,
-          quarantined = quarantined.updated(a.uid, QuarantinedTimestamp(System.nanoTime())),
+          quarantined = quarantined.updated(a.uid, QuarantinedTimestamp(System.nanoTime(), harmless)),
           _uniqueRemoteAddress)
       case None => this
     }

diff --git a/remote/src/main/scala/org/apache/pekko/remote/artery/Association.scala b/remote/src/main/scala/org/apache/pekko/remote/artery/Association.scala
@@ -538,7 +538,7 @@ private[remote] class Association(
         current.uniqueRemoteAddress() match {
           case Some(peer) if peer.uid == u =>
             if (!current.isQuarantined(u)) {
-              val newState = current.newQuarantined()
+              val newState = current.newQuarantined(harmless)
               if (swapState(current, newState)) {
                 // quarantine state change was performed
                 if (harmless) {

diff --git a/remote/src/test/scala/org/apache/pekko/remote/artery/OutboundIdleShutdownSpec.scala b/remote/src/test/scala/org/apache/pekko/remote/artery/OutboundIdleShutdownSpec.scala
@@ -117,6 +117,7 @@ class OutboundIdleShutdownSpec extends ArteryMultiNodeSpec(s"""
 
       localArtery.quarantine(remoteAddress, Some(remoteUid), "Test")
       association.associationState.isQuarantined(remoteUid) shouldBe true
+      association.associationState.quarantinedButHarmless(remoteUid) shouldBe false
 
       eventually {
         assertStreamActive(association, Association.ControlQueueIndex, expected = false)
@@ -136,6 +137,7 @@ class OutboundIdleShutdownSpec extends ArteryMultiNodeSpec(s"""
 
         localArtery.quarantine(remoteAddress, Some(remoteUid), "HarmlessTest", harmless = true)
         association.associationState.isQuarantined(remoteUid) shouldBe true
+        association.associationState.quarantinedButHarmless(remoteUid) shouldBe true
 
         eventually {
           assertStreamActive(association, Association.ControlQueueIndex, expected = false)