add self ref

Author: sebastian-alfers

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

@@ -690,7 +690,7 @@ private[akka] class ShardRegion(
     timers.startTimerWithFixedDelay(Retry, Retry, retryInterval)
     startRegistration()
     logPassivationStrategy()
-    instrumentation.shardRegionBufferSize(cluster.selfAddress, typeName, 0)
+    instrumentation.shardRegionBufferSize(cluster.selfAddress, self, typeName, 0)
   }
 
   override def postStop(): Unit = {
@@ -699,7 +699,7 @@ private[akka] class ShardRegion(
     coordinator.foreach(_ ! RegionStopped(context.self))
     cluster.unsubscribe(self)
     gracefulShutdownProgress.trySuccess(Done)
-    instrumentation.shardRegionBufferSize(cluster.selfAddress, typeName, 0)
+    instrumentation.shardRegionBufferSize(cluster.selfAddress, self, typeName, 0)
   }
 
   private def logPassivationStrategy(): Unit = {
@@ -954,7 +954,7 @@ private[akka] class ShardRegion(
             dropped,
             shard)
           // better to decrease by "dropped" to avoid calculating the size?
-          instrumentation.shardRegionBufferSize(cluster.selfAddress, typeName, shardBuffers.totalSize)
+          instrumentation.shardRegionBufferSize(cluster.selfAddress, self, typeName, shardBuffers.totalSize)
         }
         loggedFullBufferWarning = false
       }
@@ -1298,7 +1298,7 @@ private[akka] class ShardRegion(
       context.system.deadLetters ! msg
     } else {
       shardBuffers.append(shardId, msg, snd)
-      instrumentation.incrementShardRegionBufferSize(cluster.selfAddress, typeName)
+      instrumentation.incrementShardRegionBufferSize(cluster.selfAddress, self, typeName)
       // log some insight to how buffers are filled up every 10% of the buffer capacity
       val tot = totBufSize + 1
       if (tot % (bufferSize / 10) == 0) {
@@ -1331,7 +1331,7 @@ private[akka] class ShardRegion(
       }
 
       shardBuffers.remove(shardId)
-      instrumentation.shardRegionBufferSize(cluster.selfAddress, typeName, shardBuffers.totalSize)
+      instrumentation.shardRegionBufferSize(cluster.selfAddress, self, typeName, shardBuffers.totalSize)
     }
     loggedFullBufferWarning = false
     retryCount = 0
@@ -1374,7 +1374,7 @@ private[akka] class ShardRegion(
               shardId,
               buf.size + 1)
             shardBuffers.append(shardId, msg, snd)
-            instrumentation.incrementShardRegionBufferSize(cluster.selfAddress, typeName)
+            instrumentation.incrementShardRegionBufferSize(cluster.selfAddress, self, typeName)
         }
 
       case _ =>

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

@@ -9,6 +9,7 @@ import akka.annotation.InternalStableApi
 import scala.collection.immutable
 import scala.jdk.CollectionConverters._
 
+import akka.actor.ActorRef
 import akka.actor.ActorSystem
 import akka.actor.Address
 import akka.actor.ClassicActorSystemProvider
@@ -80,11 +81,18 @@ class ClusterShardingInstrumentationProvider(system: ExtendedActorSystem) extend
 class ClusterShardingTelemetryEnsemble(val instrumentations: Seq[ClusterShardingInstrumentation])
     extends ClusterShardingInstrumentation {
 
-  override def shardRegionBufferSize(selfAddress: Address, typeName: String, size: Int): Unit =
-    instrumentations.foreach(_.shardRegionBufferSize(selfAddress, typeName, size))
+  override def shardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String,
+      size: Int): Unit =
+    instrumentations.foreach(_.shardRegionBufferSize(selfAddress, shardRegionActor, typeName, size))
 
-  override def incrementShardRegionBufferSize(selfAddress: Address, typeName: String): Unit =
-    instrumentations.foreach(_.incrementShardRegionBufferSize(selfAddress, typeName))
+  override def incrementShardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String): Unit =
+    instrumentations.foreach(_.incrementShardRegionBufferSize(selfAddress, shardRegionActor, typeName))
 
   override def dependencies: immutable.Seq[String] =
     instrumentations.flatMap(_.dependencies)
@@ -102,9 +110,16 @@ object EmptyClusterShardingInstrumentation extends EmptyClusterShardingInstrumen
 @InternalStableApi
 class EmptyClusterShardingInstrumentation extends ClusterShardingInstrumentation {
 
-  override def shardRegionBufferSize(selfAddress: Address, typeName: String, size: Int): Unit = ()
+  override def shardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String,
+      size: Int): Unit = ()
 
-  override def incrementShardRegionBufferSize(selfAddress: Address, typeName: String): Unit = ()
+  override def incrementShardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String): Unit = ()
 
   override def dependencies: immutable.Seq[String] = Nil
 }
@@ -118,12 +133,12 @@ trait ClusterShardingInstrumentation {
   /**
    * @param size set current size of the buffer.
    */
-  def shardRegionBufferSize(selfAddress: Address, typeName: String, size: Int): Unit
+  def shardRegionBufferSize(selfAddress: Address, shardRegionActor: ActorRef, typeName: String, size: Int): Unit
 
   /**
    * Increase the current size of the buffer by one.
    */
-  def incrementShardRegionBufferSize(selfAddress: Address, typeName: String): Unit
+  def incrementShardRegionBufferSize(selfAddress: Address, shardRegionActor: ActorRef, typeName: String): Unit
 
   /**
    * Optional dependencies for this instrumentation.

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

@@ -5,12 +5,16 @@
 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
@@ -48,11 +52,18 @@ class SpecClusterShardingTelemetry(@nowarn("msg=never used") system: ExtendedAct
 
   val counter = new AtomicInteger(0)
 
-  override def shardRegionBufferSize(selfAddress: Address, typeName: String, size: Int): Unit = {
+  override def shardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String,
+      size: Int): Unit = {
     counter.set(size)
   }
 
-  override def incrementShardRegionBufferSize(selfAddress: Address, typeName: String): Unit = {
+  override def incrementShardRegionBufferSize(
+      selfAddress: Address,
+      shardRegionActor: ActorRef,
+      typeName: String): Unit = {
     counter.incrementAndGet()
   }