feat: latency for shard handoff

Author: sebastian-alfers

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

@@ -34,6 +34,7 @@ import akka.cluster.sharding.internal.{
 import akka.cluster.sharding.internal.AbstractLeastShardAllocationStrategy
 import akka.cluster.sharding.internal.ClusterShardAllocationMixin.RegionEntry
 import akka.cluster.sharding.internal.ClusterShardAllocationMixin.ShardSuitabilityOrdering
+import akka.cluster.sharding.internal.ClusterShardingInstrumentationProvider
 import akka.cluster.sharding.internal.EventSourcedRememberEntitiesCoordinatorStore.MigrationMarker
 import akka.event.{ BusLogging, Logging }
 import akka.pattern.{ pipe, AskTimeoutException }
@@ -644,6 +645,11 @@ object ShardCoordinator {
 
     import Internal._
 
+    private val instrumentation = ClusterShardingInstrumentationProvider.get(context.system).instrumentation
+    val cluster = Cluster(context.system)
+
+    instrumentation.beginnShardHandoff(cluster.selfAddress, self, typeName)
+
     regions.foreach { region =>
       region ! BeginHandOff(shard)
     }
@@ -707,6 +713,11 @@ object ShardCoordinator {
     }
 
     def done(ok: Boolean): Unit = {
+      if (ok) {
+        instrumentation.finishedShardHandoffOk(cluster.selfAddress, self, typeName)
+      } else {
+        instrumentation.finishedShardHandoffNotOk(cluster.selfAddress, self, typeName)
+      }
       context.parent ! RebalanceDone(shard, ok)
       context.stop(self)
     }

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

@@ -94,6 +94,18 @@ class ClusterShardingTelemetryEnsemble(val instrumentations: Seq[ClusterSharding
       typeName: String): Unit =
     instrumentations.foreach(_.incrementShardRegionBufferSize(selfAddress, shardRegionActor, typeName))
 
+  override def beginnShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit =
+    instrumentations.foreach(_.beginnShardHandoff(selfAddress, shardCoordinatorActor, typeName))
+
+  override def finishedShardHandoffOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit =
+    instrumentations.foreach(_.finishedShardHandoffOk(selfAddress, shardCoordinatorActor, typeName))
+
+  override def finishedShardHandoffNotOk(
+      selfAddress: Address,
+      shardCoordinatorActor: ActorRef,
+      typeName: String): Unit =
+    instrumentations.foreach(_.finishedShardHandoffNotOk(selfAddress, shardCoordinatorActor, typeName))
+
   override def dependencies: immutable.Seq[String] =
     instrumentations.flatMap(_.dependencies)
 }
@@ -121,6 +133,12 @@ class EmptyClusterShardingInstrumentation extends ClusterShardingInstrumentation
       shardRegionActor: ActorRef,
       typeName: String): Unit = ()
 
+  override def beginnShardHandoff(selfAddress: Address, shardRegionActor: ActorRef, typeName: String): Unit = ()
+
+  override def finishedShardHandoffOk(selfAddress: Address, self: ActorRef, typeName: String): Unit = ()
+
+  override def finishedShardHandoffNotOk(selfAddress: Address, self: ActorRef, typeName: String): Unit = ()
+
   override def dependencies: immutable.Seq[String] = Nil
 }
 
@@ -140,6 +158,12 @@ trait ClusterShardingInstrumentation {
    */
   def incrementShardRegionBufferSize(selfAddress: Address, shardRegionActor: ActorRef, typeName: String): Unit
 
+  def beginnShardHandoff(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
+
+  def finishedShardHandoffOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
+
+  def finishedShardHandoffNotOk(selfAddress: Address, shardCoordinatorActor: ActorRef, typeName: String): Unit
+
   /**
    * Optional dependencies for this instrumentation.
    *

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

@@ -4,23 +4,17 @@
 
 package akka.cluster.sharding
 
-import java.util.concurrent.atomic.AtomicInteger
-
-import scala.annotation.nowarn
-
 import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.time.{ Seconds, Span }
 import scala.concurrent.duration._
 
 import org.scalatest.concurrent.Eventually.eventually
 
-import akka.actor.ActorRef
-import akka.actor.ExtendedActorSystem
 import akka.actor.{ Actor, ActorLogging, Address, Props }
 import akka.cluster.Cluster
 import akka.cluster.MemberStatus
 import akka.cluster.sharding.ClusterShardingInstrumentationSpec.GiveMeYourHome.{ Get, Home }
-import akka.cluster.sharding.internal.{ ClusterShardingInstrumentation, ClusterShardingInstrumentationProvider }
+import akka.cluster.sharding.internal.ClusterShardingInstrumentationProvider
 import akka.remote.testkit.Direction
 import akka.testkit.TestProbe
 import akka.serialization.jackson.CborSerializable
@@ -32,7 +26,7 @@ object ClusterShardingInstrumentationSpecConfig
       additionalConfig = """
       akka.cluster.sharding {
         rebalance-interval = 120 s
-        telemetry.instrumentations += akka.cluster.sharding.SpecClusterShardingTelemetry
+        telemetry.instrumentations += akka.cluster.sharding.ClusterShardingInstrumentationSpecTelemetry
         buffer-size = 120
       }
      """) {
@@ -47,29 +41,6 @@ class ClusterShardingInstrumentationSpecMultiJvmNode1 extends ClusterShardingIns
 
 class ClusterShardingInstrumentationSpecMultiJvmNode2 extends ClusterShardingInstrumentationSpec
 
-class SpecClusterShardingTelemetry(@nowarn("msg=never used") system: ExtendedActorSystem)
-    extends ClusterShardingInstrumentation {
-
-  val counter = new AtomicInteger(0)
-
-  override def shardRegionBufferSize(
-      selfAddress: Address,
-      shardRegionActor: ActorRef,
-      typeName: String,
-      size: Int): Unit = {
-    counter.set(size)
-  }
-
-  override def incrementShardRegionBufferSize(
-      selfAddress: Address,
-      shardRegionActor: ActorRef,
-      typeName: String): Unit = {
-    counter.incrementAndGet()
-  }
-
-  override def dependencies: Seq[String] = Nil
-}
-
 object ClusterShardingInstrumentationSpec {
 
   object GiveMeYourHome {
@@ -111,7 +82,7 @@ abstract class ClusterShardingInstrumentationSpec
   import ClusterShardingInstrumentationSpecConfig._
 
   private val counter = ClusterShardingInstrumentationProvider(system).instrumentation
-    .asInstanceOf[SpecClusterShardingTelemetry]
+    .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
     .counter
 
   override implicit val patienceConfig: PatienceConfig = {

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

@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2025 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.cluster.sharding
+
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.annotation.nowarn
+
+import akka.actor.ActorRef
+import akka.actor.Address
+import akka.actor.ExtendedActorSystem
+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)
+
+  override def shardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String,
+      size: Int): Unit = {
+    counter.set(size)
+  }
+
+  override def incrementShardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String): Unit = {
+    counter.incrementAndGet()
+  }
+
+  override def beginnShardHandoff(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 finishedShardHandoffNotOk(
+      selfAddress: Address,
+      shardCoordinatorActor: ActorRef,
+      typeName: String): Unit = {
+    println("ok jea >?>> NOT FINNNNNNN")
+  }
+
+  override def dependencies: Seq[String] = Nil
+}

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

@@ -14,6 +14,7 @@ import akka.actor.Address
 import akka.actor.Props
 import akka.cluster.Cluster
 import akka.cluster.MemberStatus.Up
+import akka.cluster.sharding.internal.ClusterShardingInstrumentationProvider
 import akka.serialization.jackson.CborSerializable
 import akka.testkit.ImplicitSender
 
@@ -29,6 +30,8 @@ object RollingUpdateShardAllocationSpecConfig
         akka.coordinated-shutdown.terminate-actor-system = off
         # use the new LeastShardAllocationStrategy
         akka.cluster.sharding.least-shard-allocation-strategy.rebalance-absolute-limit = 1
+        # to test shard handoff latency telemetry
+        telemetry.instrumentations += akka.cluster.sharding.ClusterShardingInstrumentationSpecTelemetry
       }
      """) {
 
@@ -220,7 +223,16 @@ abstract class RollingUpdateShardAllocationSpec
       }
       enterBarrier("completo")
     }
+    "verify instrumentation" in {
+      // only run on oldest - Shard Coordinator
+      runOn(third) {
+        val shardHandoffDuration = ClusterShardingInstrumentationProvider(system).instrumentation
+          .asInstanceOf[ClusterShardingInstrumentationSpecTelemetry]
+          .shardHandoffDuration
 
+        shardHandoffDuration.get() > 0 shouldBe true
+      }
+      enterBarrier("shard-handoff-latency")
+    }
   }
-
 }