feat: replace CountMinSketch with FastFrequencySketch in Artery compression

remote/src/main/resources/reference.conf

@@ -1046,6 +1046,12 @@ pekko {
         # compression of common strings in remoting messages, like actor destinations, serializers etc
         compression {
 
+          # Frequency sketch implementation used for heavy hitter detection.
+          # Options:
+          #   - "count-min-sketch": Original CountMinSketch implementation (legacy, ~128KB per connection)
+          #   - "fast-frequency-sketch": FastFrequencySketch with TinyLFU aging (default, ~4KB per connection)
+          frequency-sketch-implementation = "fast-frequency-sketch"
+
           actor-refs {
             # Max number of compressed actor-refs
             # Note that compression tables are "rolling" (i.e. a new table replaces the old

remote/src/main/scala/org/apache/pekko/remote/artery/ArterySettings.scala

@@ -278,6 +278,12 @@ private[pekko] object ArterySettings {
 
     private[pekko] final val Enabled = ActorRefs.Enabled || Manifests.Enabled
 
+    val FrequencySketchImplementation: String = toRootLowerCase(getString("frequency-sketch-implementation")) match {
+      case "count-min-sketch"      => "count-min-sketch"
+      case "fast-frequency-sketch" => "fast-frequency-sketch"
+      case other                   => throw new IllegalArgumentException(s"Unknown frequency-sketch-implementation: $other")
+    }
+
     object ActorRefs {
       val config: Config = getConfig("actor-refs")
       import config._

remote/src/main/scala/org/apache/pekko/remote/artery/compress/InboundCompressions.scala

@@ -28,8 +28,54 @@ import pekko.actor.InternalActorRef
 import pekko.event.Logging
 import pekko.event.LoggingAdapter
 import pekko.remote.artery._
+import pekko.util.FastFrequencySketch
 import pekko.util.OptionVal
 
+/**
+ * INTERNAL API
+ *
+ * Abstraction over frequency sketch implementations used for heavy hitter detection in Artery compression.
+ */
+private[remote] trait FrequencySketch[T] {
+  def increment(value: T): Unit
+  def frequency(value: T): Long
+}
+
+/**
+ * INTERNAL API
+ *
+ * Wrapper around CountMinSketch that implements the FrequencySketch interface.
+ */
+private[remote] final class CountMinSketchFrequencySketch[T](depth: Int, width: Int, seed: Int)
+    extends FrequencySketch[T] {
+  private val cms = new CountMinSketch(depth, width, seed)
+
+  override def increment(value: T): Unit = {
+    cms.addObjectAndEstimateCount(value, 1)
+  }
+
+  override def frequency(value: T): Long = {
+    cms.estimateCount(value)
+  }
+}
+
+/**
+ * INTERNAL API
+ *
+ * Wrapper around FastFrequencySketch that implements the FrequencySketch interface.
+ */
+private[remote] final class FastFrequencySketchWrapper[T](capacity: Int) extends FrequencySketch[T] {
+  private val sketch = FastFrequencySketch[T](capacity)
+
+  override def increment(value: T): Unit = {
+    sketch.increment(value)
+  }
+
+  override def frequency(value: T): Long = {
+    sketch.frequency(value).toLong
+  }
+}
+
 /**
  * INTERNAL API
  * Decompress and cause compression advertisements.
@@ -352,7 +398,14 @@ private[remote] abstract class InboundCompression[T >: Null](
   private var resendCount = 0
   private val maxResendCount = 3
 
-  private val cms = new CountMinSketch(16, 1024, System.currentTimeMillis().toInt)
+  private val frequencySketch: FrequencySketch[T] = settings.FrequencySketchImplementation match {
+    case "count-min-sketch" =>
+      new CountMinSketchFrequencySketch[T](depth = 16, width = 1024, seed = System.currentTimeMillis().toInt)
+    case "fast-frequency-sketch" =>
+      new FastFrequencySketchWrapper[T](capacity = settings.ActorRefs.Max)
+    case other =>
+      throw new IllegalStateException(s"Unknown frequency-sketch-implementation: $other")
+  }
 
   log.debug("Initializing {} for originUid [{}]", Logging.simpleName(getClass), originUid)
 
@@ -442,8 +495,13 @@ private[remote] abstract class InboundCompression[T >: Null](
    * Empty keys are omitted.
    */
   def increment(@nowarn("msg=never used") remoteAddress: Address, value: T, n: Long): Unit = {
-    val count = cms.addObjectAndEstimateCount(value, n)
-    addAndCheckIfheavyHitterDetected(value, count)
+    var i = 0
+    while (i < n) {
+      frequencySketch.increment(value)
+      i += 1
+    }
+    val frequency = frequencySketch.frequency(value)
+    addAndCheckIfheavyHitterDetected(value, frequency)
     alive = true
   }
 
@@ -537,7 +595,7 @@ private[remote] abstract class InboundCompression[T >: Null](
   }
 
   override def toString =
-    s"""${Logging.simpleName(getClass)}(countMinSketch: $cms, heavyHitters: $heavyHitters)"""
+    s"""${Logging.simpleName(getClass)}(frequencySketch: $frequencySketch, heavyHitters: $heavyHitters)"""
 
 }
 

remote/src/main/scala/org/apache/pekko/remote/artery/compress/TopHeavyHitters.scala

@@ -150,50 +150,52 @@ private[remote] final class TopHeavyHitters[T >: Null](val max: Int)(implicit cl
    * Attempt adding item to heavy hitters set, if it does not fit in the top yet,
    * it will be dropped and the method will return `false`.
    *
-   * WARNING: It is only allowed to *increase* the weight of existing elements, decreasing is disallowed.
+   * Existing entries always have their weight updated (both increase and decrease are supported),
+   * which is needed when using a frequency sketch with periodic reset (aging).
    *
    * @return `true` if the added item has become a heavy hitter.
    */
   // TODO possibly can be optimised further? (there is a benchmark)
-  def update(item: T, count: Long): Boolean =
-    isHeavy(count) && { // O(1) terminate execution ASAP if known to not be a heavy hitter anyway
-      val hashCode = new HashCodeVal(item.hashCode()) // avoid re-calculating hashCode
-      val startIndex = hashCode.get & mask
-
-      // We first try to find the slot where an element with an equal hash value is. This is a possible candidate
-      // for an actual matching entry (unless it is an entry with a colliding hash value).
-      // worst case O(n), common O(1 + alpha), can't really bin search here since indexes are kept in synch with other arrays hmm...
-      val candidateIndex = findHashIdx(startIndex, hashCode)
-
-      if (candidateIndex == -1) {
-        // No matching hash value entry is found, so we are sure we don't have this entry yet.
-        insertKnownNewHeavy(hashCode, item, count) // O(log n + alpha)
-        true
+  def update(item: T, count: Long): Boolean = {
+    val hashCode = new HashCodeVal(item.hashCode()) // avoid re-calculating hashCode
+    val startIndex = hashCode.get & mask
+
+    // We first try to find the slot where an element with an equal hash value is. This is a possible candidate
+    // for an actual matching entry (unless it is an entry with a colliding hash value).
+    // worst case O(n), common O(1 + alpha), can't really bin search here since indexes are kept in synch with other arrays hmm...
+    val candidateIndex = findHashIdx(startIndex, hashCode)
+
+    if (candidateIndex == -1) {
+      // No matching hash value entry is found, so we are sure we don't have this entry yet.
+      // Only insert new entries if the weight qualifies as heavy.
+      isHeavy(count) && { insertKnownNewHeavy(hashCode, item, count); true }
+    } else {
+      // We now found, relatively cheaply, the first index where our searched entry *might* be (hashes are equal).
+      // This is not guaranteed to be the one we are searching for, yet (hash values may collide).
+      // From this position we can invoke the more costly search which checks actual object equalities.
+      // With this two step search we avoid equality checks completely for many non-colliding entries.
+      val actualIdx = findItemIdx(candidateIndex, hashCode, item)
+
+      // usually O(1), worst case O(n) if we need to scan due to hash conflicts
+      if (actualIdx == -1) {
+        // So we don't have this entry so far (only a colliding one, it was a false positive from findHashIdx).
+        // Only insert new entries if the weight qualifies as heavy.
+        isHeavy(count) && { insertKnownNewHeavy(hashCode, item, count); true }
       } else {
-        // We now found, relatively cheaply, the first index where our searched entry *might* be (hashes are equal).
-        // This is not guaranteed to be the one we are searching for, yet (hash values may collide).
-        // From this position we can invoke the more costly search which checks actual object equalities.
-        // With this two step search we avoid equality checks completely for many non-colliding entries.
-        val actualIdx = findItemIdx(candidateIndex, hashCode, item)
-
-        // usually O(1), worst case O(n) if we need to scan due to hash conflicts
-        if (actualIdx == -1) {
-          // So we don't have this entry so far (only a colliding one, it was a false positive from findHashIdx).
-          insertKnownNewHeavy(hashCode, item, count) // O(1 + log n), we simply replace the current lowest heavy hitter
-          true
-        } else {
-          // The entry exists, let's update it.
-          updateExistingHeavyHitter(actualIdx, count)
-          // not a "new" heavy hitter, since we only replaced it (so it was signaled as new once before)
-          false
-        }
+        // The entry exists, let's update it.
+        // Existing heavy hitters always get their weight updated (even if decreased),
+        // which is needed when using a frequency sketch with periodic reset (aging).
+        updateExistingHeavyHitter(actualIdx, count)
+        // not a "new" heavy hitter, since we only replaced it (so it was signaled as new once before)
+        false
       }
-
     }
+  }
 
   /**
    * Checks the lowest weight entry in this structure and returns true if the given count is larger than that. In
    * other words this checks if a new entry can be added as it is larger than the known least weight.
+   * Note: this only gates insertion of new entries; existing entries always get their weight updated.
    */
   private def isHeavy(count: Long): Boolean =
     count > lowestHitterWeight
@@ -233,15 +235,18 @@ private[remote] final class TopHeavyHitters[T >: Null](val max: Int)(implicit cl
   /**
    * Replace existing heavy hitter – give it a new `count` value. This will also restore the heap property, so this
    * might make a previously lowest hitter no longer be one.
+   *
+   * Supports both weight increase and decrease. When using a frequency sketch with periodic reset (aging),
+   * the estimated frequency can decrease, so the weight must be allowed to go down as well.
    */
   private def updateExistingHeavyHitter(foundHashIndex: Int, count: Long): Unit = {
-    if (weights(foundHashIndex) > count)
-      throw new IllegalArgumentException(
-        s"Weights can be only incremented or kept the same, not decremented. " +
-        s"Previous weight was [${weights(foundHashIndex)}], attempted to modify it to [$count].")
+    val oldCount = weights(foundHashIndex)
     weights(foundHashIndex) = count // we don't need to change `hashCode`, `heapIndex` or `item`, those remain the same
-    // Position in the heap might have changed as count was incremented
-    fixHeap(heapIndex(foundHashIndex))
+    // Position in the heap might have changed as count was updated
+    if (count > oldCount)
+      fixHeap(heapIndex(foundHashIndex)) // weight increased: push down towards children
+    else if (count < oldCount)
+      fixHeapUp(heapIndex(foundHashIndex)) // weight decreased: bubble up towards parent
   }
 
   /**
@@ -311,6 +316,23 @@ private[remote] final class TopHeavyHitters[T >: Null](val max: Int)(implicit cl
     }
   }
 
+  /**
+   * Call this if the weight of an entry at heap node `index` was decremented. Since the weight decreased,
+   * the element may now be smaller than its parent, so we need to restore the heap property "upwards".
+   */
+  @tailrec
+  private def fixHeapUp(index: Int): Unit = {
+    if (index > 0) {
+      val parentIndex = (index - 1) / 2
+      val currentWeight: Long = weights(heap(index))
+      val parentWeight: Long = weights(heap(parentIndex))
+      if (currentWeight < parentWeight) {
+        swapHeapNode(index, parentIndex)
+        fixHeapUp(parentIndex)
+      }
+    }
+  }
+
   /**
    * Swaps two elements in `heap` array and maintain correct index in `heapIndex`.
    *

remote/src/test/scala/org/apache/pekko/remote/artery/compress/CompressionIntegrationSpec.scala

@@ -146,6 +146,54 @@ class CompressionIntegrationSpec
       }
     }
 
+    "be advertised with count-min-sketch frequency sketch implementation" in {
+      val config = """
+       pekko.remote.artery.advanced.compression {
+         frequency-sketch-implementation = "count-min-sketch"
+         actor-refs.advertisement-interval = 2 seconds
+         manifests.advertisement-interval = 2 seconds
+       }
+      """
+      val systemC = newRemoteSystem(Some(config), name = Some("systemC"))
+      val systemD = newRemoteSystem(Some(config), name = Some("systemD"))
+
+      val cRefProbe = TestProbe()(systemC)
+      val dRefProbe = TestProbe()(systemD)
+      systemC.eventStream
+        .subscribe(cRefProbe.ref, classOf[CompressionProtocol.Events.ReceivedActorRefCompressionTable])
+      systemD.eventStream
+        .subscribe(dRefProbe.ref, classOf[CompressionProtocol.Events.ReceivedActorRefCompressionTable])
+
+      val echoRefD = systemD.actorOf(TestActors.echoActorProps, "echo-cms")
+
+      val cProbe = TestProbe()(systemC)
+      systemC.actorSelection(rootActorPath(systemD) / "user" / "echo-cms").tell(Identify(None), cProbe.ref)
+      val echoRefC = cProbe.expectMsgType[ActorIdentity].ref.get
+
+      (1 to messagesToExchange).foreach { _ =>
+        echoRefC.tell(TestMessage("hello"), cProbe.ref)
+      }
+      cProbe.receiveN(messagesToExchange)
+
+      within(10.seconds) {
+        awaitAssert {
+          val c1 = cRefProbe.expectMsgType[Events.ReceivedActorRefCompressionTable](2.seconds)
+          info("System [C] received: " + c1)
+          c1.table.version.toInt should be >= 1
+          c1.table.dictionary.keySet should contain(echoRefC)
+        }
+        awaitAssert {
+          val d1 = dRefProbe.expectMsgType[Events.ReceivedActorRefCompressionTable](2.seconds)
+          info("System [D] received: " + d1)
+          d1.table.version.toInt should be >= 1
+          d1.table.dictionary.keySet should contain(echoRefD)
+        }
+      }
+
+      shutdown(systemC)
+      shutdown(systemD)
+    }
+
     "not be advertised if ActorRef compression disabled" in {
       val config = """
        pekko.remote.artery.advanced.compression.actor-refs.max = off

remote/src/test/scala/org/apache/pekko/remote/artery/compress/HeavyHittersSpec.scala

@@ -170,6 +170,40 @@ class HeavyHittersSpec extends AnyWordSpecLike with Matchers {
       hitters.lowestHitterWeight should ===(3)
     }
 
+    "support weight decrease (e.g. from frequency sketch reset)" in {
+      val hitters = new TopHeavyHitters[String](2)
+      hitters.update("A", 10) should ===(true)
+      hitters.update("B", 20) should ===(true)
+      hitters.lowestHitterWeight should ===(10)
+
+      // Simulate frequency sketch reset: weight decreases
+      hitters.update("A", 5) should ===(false)
+      // A is now the lowest hitter
+      hitters.lowestHitterWeight should ===(5)
+
+      // A new item with weight > 5 should replace A
+      hitters.update("C", 8) should ===(true)
+      hitters.iterator.toSet should ===(Set("B", "C"))
+    }
+
+    "restore heap property upward when weight decreases" in {
+      val hitters = new TopHeavyHitters[String](4)
+      hitters.update("A", 10) should ===(true)
+      hitters.update("B", 20) should ===(true)
+      hitters.update("C", 30) should ===(true)
+      hitters.update("D", 40) should ===(true)
+      hitters.lowestHitterWeight should ===(10)
+
+      // Decrease D's weight significantly - it should bubble up to become the lowest
+      hitters.update("D", 1) should ===(false)
+      hitters.lowestHitterWeight should ===(1)
+      hitters.iterator.toSet should ===(Set("A", "B", "C", "D"))
+
+      // A new item should replace D (the new lowest)
+      hitters.update("E", 15) should ===(true)
+      hitters.iterator.toSet should ===(Set("A", "B", "C", "E"))
+    }
+
     "be disabled with max=0" in {
       val hitters = new TopHeavyHitters[String](0)
       hitters.update("A", 10) shouldBe true