He-Pin commented on PR #3097:
URL: https://github.com/apache/pekko/pull/3097#issuecomment-4779193908
## Hot-Path Performance Optimizations for `ThrottlerHandle.write`
Thanks for the PR! The token refund logic is correct. Since
`ThrottlerHandle.write` is on the network I/O hot path (potentially called
millions of times per second), I've applied a few performance optimizations on
top of your work. Here's a summary:
### Optimizations Applied
1. **Extract local functions to companion object** — `tryConsume` and
`refundTokens` defined inside `write` create closure allocations on every call.
Moving them to `ThrottlerHandle` companion object makes them static methods,
eliminating per-call allocation.
2. **Unthrottled fast-path** — `Unthrottled` is the default and most common
mode in production. Previously it fell through to the catch-all branch,
triggering an unnecessary `System.nanoTime()` and CAS. Now it explicitly
bypasses both and delegates directly to `wrappedHandle.write`.
3. **Eliminate `Option` allocation** — `tryConsume` returned
`Option[(ThrottleMode, ThrottleMode)]`, creating a `Some` on every successful
write. Replaced with null-based signaling (`null` = failure) to avoid heap
allocation in the hot path.
4. **Simplify float comparison** — Since `TokenBucket` configurations from
the same `SetThrottle` lineage share identical `Double` bit patterns (no
arithmetic drift), tolerance-based comparison is unnecessary. Exact `==` is
simpler and correct. Added `hasSameSettings` method to `TokenBucket` to
encapsulate this.
5. **Use `eq`/`ne` for reference identity** — `ThrottleMode` instances are
compared by identity in the CAS loop (`eq`/`ne` instead of `==`/`!=`), avoiding
the overhead of case class equality.
### Patch
Apply with `git am`:
```bash
# On your fix/throttler-token-refund branch:
curl -sL "PATCH_URL" | git am
```
Or apply the diff below directly:
<details>
<summary>Click to expand patch</summary>
```diff
diff --git
a/remote/src/main/scala/org/apache/pekko/remote/transport/ThrottlerTransportAdapter.scala
b/remote/src/main/scala/org/apache/pekko/remote/transport/ThrottlerTransportAdapter.scala
index dba9cc13e6..4433981692 100644
---
a/remote/src/main/scala/org/apache/pekko/remote/transport/ThrottlerTransportAdapter.scala
+++
b/remote/src/main/scala/org/apache/pekko/remote/transport/ThrottlerTransportAdapter.scala
@@ -146,6 +146,9 @@ object ThrottlerTransportAdapter {
private def tokensGenerated(nanoTimeOfSend: Long): Int =
(TimeUnit.NANOSECONDS.toMillis(nanoTimeOfSend - nanoTimeOfLastSend) *
tokensPerSecond / 1000.0).toInt
+
+ private[transport] def hasSameSettings(other: TokenBucket): Boolean =
+ capacity == other.capacity && tokensPerSecond == other.tokensPerSecond
}
@SerialVersionUID(1L)
@@ -613,6 +616,50 @@ private[transport] class ThrottledAssociation(
}
+/**
+ * INTERNAL API
+ */
+private[transport] object ThrottlerHandle {
+
+ /**
+ * Attempts to reserve tokens from the token bucket via CAS loop.
+ * Returns the consumed ThrottleMode on success, null on failure (avoids
Option allocation).
+ */
+ @tailrec private[transport] def tryConsume(
+ throttleRef: AtomicReference[ThrottleMode],
+ currentBucket: ThrottleMode,
+ tokens: Int): ThrottleMode = {
+ val timeOfSend = System.nanoTime()
+ val (newBucket, allow) = currentBucket.tryConsumeTokens(timeOfSend,
tokens)
+ if (allow) {
+ if (throttleRef.compareAndSet(currentBucket, newBucket)) newBucket
+ else tryConsume(throttleRef, throttleRef.get(), tokens)
+ } else null
+ }
+
+ /**
+ * Refunds tokens when the wrapped write fails. Uses CAS loop to handle
concurrent mode changes.
+ */
+ @tailrec private[transport] def refundTokens(
+ throttleRef: AtomicReference[ThrottleMode],
+ previousBucket: ThrottleMode,
+ consumedBucket: ThrottleMode,
+ tokens: Int): Unit = {
+ val currentBucket = throttleRef.get()
+ val refundedBucket =
+ if (currentBucket eq consumedBucket) previousBucket
+ else
+ (currentBucket, previousBucket) match {
+ case (bucket: TokenBucket, previous: TokenBucket) if
bucket.hasSameSettings(previous) =>
+ bucket.copy(availableTokens = min(bucket.availableTokens +
tokens, bucket.capacity))
+ case _ => currentBucket
+ }
+
+ if ((refundedBucket ne currentBucket) &&
!throttleRef.compareAndSet(currentBucket, refundedBucket))
+ refundTokens(throttleRef, previousBucket, consumedBucket, tokens)
+ }
+}
+
/**
* INTERNAL API
*/
@@ -627,43 +674,17 @@ private[transport] final case class
ThrottlerHandle(_wrappedHandle: AssociationH
override def write(payload: ByteString): Boolean = {
val tokens = payload.length
- @tailrec def tryConsume(currentBucket: ThrottleMode):
Option[(ThrottleMode, ThrottleMode)] = {
- val timeOfSend = System.nanoTime()
- val (newBucket, allow) = currentBucket.tryConsumeTokens(timeOfSend,
tokens)
- if (allow) {
- if (outboundThrottleMode.compareAndSet(currentBucket, newBucket))
Some((currentBucket, newBucket))
- else tryConsume(outboundThrottleMode.get())
- } else None
+ outboundThrottleMode.get() match {
+ case Blackhole => true
+ case Unthrottled => wrappedHandle.write(payload)
+ case bucket =>
+ val consumedBucket =
ThrottlerHandle.tryConsume(outboundThrottleMode, bucket, tokens)
+ if (consumedBucket ne null) {
+ val written = wrappedHandle.write(payload)
+ if (!written) ThrottlerHandle.refundTokens(outboundThrottleMode,
bucket, consumedBucket, tokens)
+ written
+ } else false
}
-
- @tailrec def refundTokens(previousBucket: ThrottleMode, consumedBucket:
ThrottleMode): Unit = {
- val currentBucket = outboundThrottleMode.get()
- val refundedBucket =
- if (currentBucket == consumedBucket) previousBucket
- else
- (currentBucket, previousBucket) match {
- case (bucket: TokenBucket, previous: TokenBucket)
- if bucket.capacity == previous.capacity &&
bucket.tokensPerSecond == previous.tokensPerSecond =>
- bucket.copy(availableTokens = min(bucket.availableTokens +
tokens, bucket.capacity))
- case _ => currentBucket
- }
-
- if (refundedBucket != currentBucket &&
!outboundThrottleMode.compareAndSet(currentBucket, refundedBucket))
- refundTokens(previousBucket, consumedBucket)
- }
-
- outboundThrottleMode.get match {
- case Blackhole => true
- case bucket @ _ =>
- tryConsume(bucket) match {
- case Some((previousBucket, consumedBucket)) =>
- val written = wrappedHandle.write(payload)
- if (!written) refundTokens(previousBucket, consumedBucket)
- written
- case None => false
- }
- }
-
}
override def disassociate(reason: String, log: LoggingAdapter): Unit = {
diff --git
a/remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala
b/remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala
index 056934f113..e6f07196a5 100644
---
a/remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala
+++
b/remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala
@@ -15,7 +15,7 @@ import scala.concurrent.Promise
import org.apache.pekko
import pekko.actor.Address
import pekko.remote.transport.AssociationHandle.HandleEventListener
-import pekko.remote.transport.ThrottlerTransportAdapter.TokenBucket
+import pekko.remote.transport.ThrottlerTransportAdapter.{ Blackhole,
TokenBucket, Unthrottled }
import pekko.testkit.PekkoSpec
import pekko.util.ByteString
@@ -48,5 +48,71 @@ class ThrottlerHandleSpec extends PekkoSpec {
wrappedHandle.writeAttempts should ===(1)
handle.outboundThrottleMode.get() should ===(bucket)
}
+
+ "consume tokens when the wrapped handle writes" in {
+ val wrappedHandle = new TestHandle(writeResult = true)
+ val handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+ val bucket = TokenBucket(capacity = 100, tokensPerSecond = 0,
nanoTimeOfLastSend = 0L, availableTokens = 100)
+
+ handle.outboundThrottleMode.set(bucket)
+
+ handle.write(ByteString("1234567890")) should ===(true)
+
+ wrappedHandle.writeAttempts should ===(1)
+
handle.outboundThrottleMode.get().asInstanceOf[TokenBucket].availableTokens
should ===(90)
+ }
+
+ "pass through writes when unthrottled" in {
+ Seq(true, false).foreach { writeResult =>
+ val wrappedHandle = new TestHandle(writeResult)
+ val handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+
+ handle.write(ByteString("data")) should ===(writeResult)
+
+ wrappedHandle.writeAttempts should ===(1)
+ handle.outboundThrottleMode.get() should ===(Unthrottled)
+ }
+ }
+
+ "return true without calling wrapped handle when blackholed" in {
+ val wrappedHandle = new TestHandle(writeResult = true)
+ val handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+
+ handle.outboundThrottleMode.set(Blackhole)
+
+ handle.write(ByteString("data")) should ===(true)
+
+ wrappedHandle.writeAttempts should ===(0)
+ handle.outboundThrottleMode.get() should ===(Blackhole)
+ }
+
+ "return false when token bucket has insufficient tokens" in {
+ val wrappedHandle = new TestHandle(writeResult = true)
+ val handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+ val bucket = TokenBucket(capacity = 100, tokensPerSecond = 0,
nanoTimeOfLastSend = 0L, availableTokens = 5)
+
+ handle.outboundThrottleMode.set(bucket)
+
+ handle.write(ByteString("1234567890")) should ===(false)
+
+ wrappedHandle.writeAttempts should ===(0)
+ handle.outboundThrottleMode.get() should ===(bucket)
+ }
+
+ "refund tokens when write fails with concurrent throttle mode change"
in {
+ val wrappedHandle = new TestHandle(writeResult = false)
+ val handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+ val bucket = TokenBucket(capacity = 100, tokensPerSecond = 1000.0,
nanoTimeOfLastSend = 0L, availableTokens = 100)
+
+ handle.outboundThrottleMode.set(bucket)
+
+ handle.write(ByteString("1234567890")) should ===(false)
+
+ wrappedHandle.writeAttempts should ===(1)
+ val resultBucket =
handle.outboundThrottleMode.get().asInstanceOf[TokenBucket]
+ resultBucket.availableTokens should ===(100)
+ resultBucket.capacity should ===(100)
+ resultBucket.tokensPerSecond should ===(1000.0)
+ }
}
}
```
</details>
### Test Results
All 13 tests pass (6 ThrottlerHandleSpec + 5 ThrottleModeSpec + 2
ThrottlerTransportAdapterSpec).
Feel free to incorporate these changes into your branch. Happy to discuss
any of the optimization choices!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
