fix spec

Author: sebastian-alfers

akka-cluster-sharding/src/multi-jvm/scala/akka/cluster/sharding/ClusterShardingInstrumentationSpec.scala

@@ -94,11 +94,6 @@ abstract class ClusterShardingInstrumentationSpec
       .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
       .shardHomeRequests
 
-  val shardHomeResponses =
-    ClusterShardingInstrumentationProvider(system).instrumentation
-      .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
-      .shardHomeResponses
-
   override implicit val patienceConfig: PatienceConfig = {
     import akka.testkit.TestDuration
     PatienceConfig(testKitSettings.DefaultTimeout.duration.dilated, Span(1000, org.scalatest.time.Millis))
@@ -191,7 +186,21 @@ abstract class ClusterShardingInstrumentationSpec
 
     "record latency of requesting ShardHome" in {
       runOn(second) {
-        shardHomeRequests.get() shouldBe shardHomeResponses.get()
+        eventually(timeout(Span(5, Seconds))) {
+          shardHomeRequests.size > 0 shouldBe true
+          shardHomeRequests.foreach {
+            case (key, value) =>
+              if (key.startsWith("id")) {
+                // The "id-0", "id-1" ... messages were send during the blackhole.
+                // This means they were requested twice, but only received back once.
+                value.get() shouldBe 1
+              } else {
+                // The "a", "b", "c" messages were send before the blackhole.
+                // This means they got a proper send-and-request cycle
+                value.get() shouldBe 0
+              }
+          }
+        }
       }
       enterBarrier("measure-shard-home-latency")
     }

akka-cluster-sharding/src/multi-jvm/scala/akka/cluster/sharding/ClusterShardingInstrumentationSpecTelemetry.scala

@@ -7,6 +7,7 @@ package akka.cluster.sharding
 import java.util.concurrent.atomic.AtomicInteger
 
 import scala.annotation.nowarn
+import scala.collection.concurrent.TrieMap
 
 import akka.actor.ActorRef
 import akka.actor.Address
@@ -19,8 +20,9 @@ class ClusterShardingInstrumentationSpecTelemetry(@nowarn("msg=never used") syst
   val shardRegionBufferSizeCounter = new AtomicInteger(0)
   val beginShardHandoffDurationCounter = new AtomicInteger(0)
   val finishShardHandoffDurationCounter = new AtomicInteger(0)
-  val shardHomeRequests = new AtomicInteger(0)
-  val shardHomeResponses = new AtomicInteger(0)
+
+  val shardHomeRequests = new TrieMap[String, AtomicInteger]()
+
   val dropMessageCounter = new AtomicInteger(0)
 
   override def shardRegionBufferSize(
@@ -55,14 +57,18 @@ class ClusterShardingInstrumentationSpecTelemetry(@nowarn("msg=never used") syst
       selfAddress: Address,
       shardRegionActor: ActorRef,
       typeName: String,
-      shardId: String): Unit =
-    shardHomeRequests.incrementAndGet()
+      shardId: String): Unit = shardHomeRequests.getOrElseUpdate(shardId, new AtomicInteger(0)).incrementAndGet()
 
   override def receivedShardHome(
       selfAddress: Address,
       shardRegionActor: ActorRef,
       typeName: String,
-      shardId: String): Unit = shardHomeResponses.incrementAndGet()
+      shardId: String): Unit =
+    shardHomeRequests.get(shardId) match {
+      case Some(value) =>
+        value.decrementAndGet()
+      case None => () // should not happen
+    }
 
   override def messageDropped(selfAddress: Address, self: ActorRef, typeName: String): Unit =
     dropMessageCounter.incrementAndGet()