more tweaks based on pr review

Author: sebastian-alfers

akka-cluster-sharding/src/main/scala/akka/cluster/sharding/ShardCoordinator.scala

@@ -648,7 +648,7 @@ object ShardCoordinator {
     private val instrumentation = ClusterShardingInstrumentationProvider.get(context.system).instrumentation
     private val cluster = Cluster(context.system)
 
-    instrumentation.beginShardHandoff(cluster.selfAddress, self, typeName)
+    instrumentation.beginShardHandoff(cluster.selfAddress, self, typeName, shard)
 
     regions.foreach { region =>
       region ! BeginHandOff(shard)
@@ -713,11 +713,7 @@ object ShardCoordinator {
     }
 
     def done(ok: Boolean): Unit = {
-      if (ok) {
-        instrumentation.finishedShardHandoffOk(cluster.selfAddress, self, typeName)
-      } else {
-        instrumentation.finishedShardHandoffNotOk(cluster.selfAddress, self, typeName)
-      }
+      instrumentation.finishedShardHandoff(cluster.selfAddress, self, typeName, shard, ok)
       context.parent ! RebalanceDone(shard, ok)
       context.stop(self)
     }

akka-cluster-sharding/src/main/scala/akka/cluster/sharding/internal/ClusterShardingInstrumentation.scala

@@ -94,17 +94,20 @@ class ClusterShardingTelemetryEnsemble(val instrumentations: Seq[ClusterSharding
       typeName: String): Unit =
     instrumentations.foreach(_.incrementShardRegionBufferSize(selfAddress, shardRegionActor, typeName))
 
-  override def beginShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit =
-    instrumentations.foreach(_.beginShardHandoff(selfAddress, shardCoordinatorActor, typeName))
-
-  override def finishedShardHandoffOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit =
-    instrumentations.foreach(_.finishedShardHandoffOk(selfAddress, shardCoordinatorActor, typeName))
+  override def beginShardHandoff(
+      selfAddress: Address,
+      shardCoordinatorActor: ActorRef,
+      typeName: String,
+      shard: String): Unit =
+    instrumentations.foreach(_.beginShardHandoff(selfAddress, shardCoordinatorActor, typeName, shard))
 
-  override def finishedShardHandoffNotOk(
+  override def finishedShardHandoff(
       selfAddress: Address,
       shardCoordinatorActor: ActorRef,
-      typeName: String): Unit =
-    instrumentations.foreach(_.finishedShardHandoffNotOk(selfAddress, shardCoordinatorActor, typeName))
+      typeName: String,
+      shard: String,
+      ok: Boolean): Unit =
+    instrumentations.foreach(_.finishedShardHandoff(selfAddress, shardCoordinatorActor, typeName, shard, ok))
 
   override def dependencies: immutable.Seq[String] =
     instrumentations.flatMap(_.dependencies)
@@ -133,11 +136,18 @@ class EmptyClusterShardingInstrumentation extends ClusterShardingInstrumentation
       shardRegionActor: ActorRef,
       typeName: String): Unit = ()
 
-  override def beginShardHandoff(selfAddress: Address, shardRegionActor: ActorRef, typeName: String): Unit = ()
-
-  override def finishedShardHandoffOk(selfAddress: Address, self: ActorRef, typeName: String): Unit = ()
+  override def beginShardHandoff(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String,
+      shard: String): Unit = ()
 
-  override def finishedShardHandoffNotOk(selfAddress: Address, self: ActorRef, typeName: String): Unit = ()
+  override def finishedShardHandoff(
+      selfAddress: Address,
+      self: ActorRef,
+      typeName: String,
+      shard: String,
+      ok: Boolean): Unit = ()
 
   override def dependencies: immutable.Seq[String] = Nil
 }
@@ -158,11 +168,14 @@ trait ClusterShardingInstrumentation {
    */
   def incrementShardRegionBufferSize(selfAddress: Address, shardRegionActor: ActorRef, typeName: String): Unit
 
-  def beginShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
+  def beginShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String, shard: String): Unit
 
-  def finishedShardHandoffOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
-
-  def finishedShardHandoffNotOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
+  def finishedShardHandoff(
+      selfAddress: Address,
+      shardCoordinatorActor: ActorRef,
+      typeName: String,
+      shard: String,
+      ok: Boolean): Unit
 
   /**
    * Optional dependencies for this instrumentation.

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

@@ -83,7 +83,7 @@ abstract class ClusterShardingInstrumentationSpec
 
   private val counter = ClusterShardingInstrumentationProvider(system).instrumentation
     .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
-    .counter
+    .shardRegionBufferSizeCounter
 
   override implicit val patienceConfig: PatienceConfig = {
     import akka.testkit.TestDuration

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

@@ -5,7 +5,6 @@
 package akka.cluster.sharding
 
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.atomic.AtomicLong
 
 import scala.annotation.nowarn
 
@@ -17,43 +16,37 @@ import akka.cluster.sharding.internal.ClusterShardingInstrumentation
 class ClusterShardingInstrumentationSpecTelemetry(@nowarn("msg=never used") system: ExtendedActorSystem)
     extends ClusterShardingInstrumentation {
 
-  val counter = new AtomicInteger(0)
-  val shardHandoffDuration = new AtomicLong(0)
+  val shardRegionBufferSizeCounter = new AtomicInteger(0)
+  val beginShardHandoffDurationCounter = new AtomicInteger(0)
+  val finishShardHandoffDurationCounter = new AtomicInteger(0)
 
   override def shardRegionBufferSize(
       selfAddress: Address,
       shardRegionActor: ActorRef,
       typeName: String,
       size: Int): Unit = {
-    counter.set(size)
+    shardRegionBufferSizeCounter.set(size)
   }
 
   override def incrementShardRegionBufferSize(
       selfAddress: Address,
       shardRegionActor: ActorRef,
       typeName: String): Unit = {
-    counter.incrementAndGet()
+    shardRegionBufferSizeCounter.incrementAndGet()
   }
 
-  override def beginShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit = {
-    println("ok jea >?>>")
-    shardHandoffDuration.set(System.currentTimeMillis())
-    println(s"ok jea >?>> ${shardHandoffDuration.get()}")
-  }
-
-  override def finishedShardHandoffOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit = {
-    println("ok jea >?>> FINNNNNNN")
-    val currentValue = shardHandoffDuration.get()
-    shardHandoffDuration.set(System.currentTimeMillis() - currentValue)
-    println(s"ok jea >?>> FINNNNNNN ${shardHandoffDuration.get()}")
-  }
+  override def beginShardHandoff(
+      selfAddress: Address,
+      shardCoordinatorActor: ActorRef,
+      typeName: String,
+      shard: String): Unit = beginShardHandoffDurationCounter.incrementAndGet()
 
-  override def finishedShardHandoffNotOk(
+  override def finishedShardHandoff(
       selfAddress: Address,
       shardCoordinatorActor: ActorRef,
-      typeName: String): Unit = {
-    println("ok jea >?>> NOT FINNNNNNN")
-  }
+      typeName: String,
+      shard: String,
+      ok: Boolean): Unit = finishShardHandoffDurationCounter.incrementAndGet()
 
   override def dependencies: Seq[String] = Nil
 }

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

@@ -226,11 +226,14 @@ abstract class RollingUpdateShardAllocationSpec
     "verify instrumentation" in {
       // only run on oldest - Shard Coordinator
       runOn(third) {
-        val shardHandoffDuration = ClusterShardingInstrumentationProvider(system).instrumentation
+        val beginShardHandoffDurationCounter = ClusterShardingInstrumentationProvider(system).instrumentation
           .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
-          .shardHandoffDuration
+          .beginShardHandoffDurationCounter
+        val finishShardHandoffDurationCounter = ClusterShardingInstrumentationProvider(system).instrumentation
+          .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
+          .finishShardHandoffDurationCounter
 
-        shardHandoffDuration.get() > 0 shouldBe true
+        beginShardHandoffDurationCounter.get() shouldBe finishShardHandoffDurationCounter.get()
       }
       enterBarrier("shard-handoff-latency")
     }