diff --git a/remote/src/main/resources/reference.conf b/remote/src/main/resources/reference.conf index b4bf9bac8e1..ab71456d3c7 100644 --- a/remote/src/main/resources/reference.conf +++ b/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 diff --git a/remote/src/main/scala/org/apache/pekko/remote/artery/ArterySettings.scala b/remote/src/main/scala/org/apache/pekko/remote/artery/ArterySettings.scala index cb3ef30a9db..870f4057813 100644 --- a/remote/src/main/scala/org/apache/pekko/remote/artery/ArterySettings.scala +++ b/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._ diff --git a/remote/src/main/scala/org/apache/pekko/remote/artery/compress/InboundCompressions.scala b/remote/src/main/scala/org/apache/pekko/remote/artery/compress/InboundCompressions.scala index dbb43678a84..a97c28c08ec 100644 --- a/remote/src/main/scala/org/apache/pekko/remote/artery/compress/InboundCompressions.scala +++ b/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)""" } diff --git a/remote/src/main/scala/org/apache/pekko/remote/artery/compress/TopHeavyHitters.scala b/remote/src/main/scala/org/apache/pekko/remote/artery/compress/TopHeavyHitters.scala index c5458c9ed63..567f5e8476b 100644 --- a/remote/src/main/scala/org/apache/pekko/remote/artery/compress/TopHeavyHitters.scala +++ b/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`. * diff --git a/remote/src/test/scala/org/apache/pekko/remote/artery/compress/CompressionIntegrationSpec.scala b/remote/src/test/scala/org/apache/pekko/remote/artery/compress/CompressionIntegrationSpec.scala index 84c50cf0af4..834c6d6065a 100644 --- a/remote/src/test/scala/org/apache/pekko/remote/artery/compress/CompressionIntegrationSpec.scala +++ b/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 diff --git a/remote/src/test/scala/org/apache/pekko/remote/artery/compress/HeavyHittersSpec.scala b/remote/src/test/scala/org/apache/pekko/remote/artery/compress/HeavyHittersSpec.scala index 5c61ebe9e4b..c1128e7f41b 100644 --- a/remote/src/test/scala/org/apache/pekko/remote/artery/compress/HeavyHittersSpec.scala +++ b/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