This is an automated email from the ASF dual-hosted git repository.
He-Pin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/pekko.git
The following commit(s) were added to refs/heads/main by this push:
new f3654268e7 Preserve throttler tokens on failed writes (#3097)
f3654268e7 is described below
commit f3654268e763be9b77024ad13d6924ae0c483376
Author: Minh Vu <[email protected]>
AuthorDate: Tue Jun 23 20:37:03 2026 +0200
Preserve throttler tokens on failed writes (#3097)
Motivation:
ThrottlerHandle reserved token-bucket capacity before delegating to the
wrapped transport handle. Failed wrapped writes left capacity depleted even
though no payload was sent.
Modification:
Move token reservation and refund helpers out of the write hot path, add an
explicit unthrottled fast path, avoid allocation on successful reservations,
and refund failed writes without disturbing blackhole or successful-write
behavior. Add regression coverage for successful writes, unthrottled writes,
blackhole writes, insufficient tokens, and refund after a same-settings
throttle state change.
Result:
Failed wrapped writes preserve outbound throttle capacity while common
successful and unthrottled writes avoid the extra allocation path.
Tests:
- sbt "remote / Test / testOnly
org.apache.pekko.remote.transport.ThrottlerHandleSpec
org.apache.pekko.remote.classic.transport.ThrottleModeSpec
org.apache.pekko.remote.classic.transport.ThrottlerTransportAdapterSpec": passed
- sbt "remote / scalafmtAll" "remote / headerCheck": passed
- git diff --check: passed
References:
Refs #2825
---
.../transport/ThrottlerTransportAdapter.scala | 63 +++++++---
.../remote/transport/ThrottlerHandleSpec.scala | 130 +++++++++++++++++++++
2 files changed, 179 insertions(+), 14 deletions(-)
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 2660b21b19..37afab3d90 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,43 @@ private[transport] class ThrottledAssociation(
}
+/**
+ * INTERNAL API
+ */
+private[transport] object ThrottlerHandle {
+
+ @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
+ }
+
+ @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,21 +667,16 @@ private[transport] final case class
ThrottlerHandle(_wrappedHandle: AssociationH
override def write(payload: ByteString): Boolean = {
val tokens = payload.length
- @tailrec def tryConsume(currentBucket: ThrottleMode): Boolean = {
- val timeOfSend = System.nanoTime()
- val (newBucket, allow) = currentBucket.tryConsumeTokens(timeOfSend,
tokens)
- if (allow) {
- if (outboundThrottleMode.compareAndSet(currentBucket, newBucket)) true
- else tryConsume(outboundThrottleMode.get())
- } else false
- }
-
outboundThrottleMode.get match {
- case Blackhole => true
- case bucket @ _ =>
- val success = tryConsume(outboundThrottleMode.get())
- if (success) wrappedHandle.write(payload) else false
- // FIXME: this depletes the token bucket even when no write happened!!
See #2825
+ 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
}
}
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
new file mode 100644
index 0000000000..805e14ac65
--- /dev/null
+++
b/remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.remote.transport
+
+import scala.annotation.nowarn
+import scala.concurrent.Promise
+
+import org.apache.pekko
+import pekko.actor.Address
+import pekko.remote.transport.AssociationHandle.HandleEventListener
+import pekko.remote.transport.ThrottlerTransportAdapter.{ Blackhole,
TokenBucket, Unthrottled }
+import pekko.testkit.PekkoSpec
+import pekko.util.ByteString
+
+@nowarn("msg=deprecated")
+class ThrottlerHandleSpec extends PekkoSpec {
+
+ class TestHandle(writeResult: Boolean, onWrite: () => Unit = () => ())
extends AssociationHandle {
+ var writeAttempts = 0
+
+ override val localAddress: Address = Address("pekko", "local")
+ override val remoteAddress: Address = Address("pekko", "remote")
+ override val readHandlerPromise: Promise[HandleEventListener] = Promise()
+
+ override def write(payload: ByteString): Boolean = {
+ writeAttempts += 1
+ onWrite()
+ writeResult
+ }
+ }
+
+ "A ThrottlerHandle" must {
+ "preserve token bucket state when the wrapped handle does not write" in {
+ val wrappedHandle = new TestHandle(writeResult = false)
+ 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 ===(false)
+
+ 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 throttle state changes before a failed write is
refunded" in {
+ var handle: ThrottlerHandle = null
+ val bucket = TokenBucket(capacity = 100, tokensPerSecond = 0,
nanoTimeOfLastSend = 0L, availableTokens = 100)
+ val wrappedHandle = new TestHandle(
+ writeResult = false,
+ onWrite = () =>
handle.outboundThrottleMode.set(bucket.copy(nanoTimeOfLastSend = 1L,
availableTokens = 40)))
+
+ handle = ThrottlerHandle(wrappedHandle, system.deadLetters)
+ handle.outboundThrottleMode.set(bucket)
+
+ handle.write(ByteString("1234567890")) should ===(false)
+
+ wrappedHandle.writeAttempts should ===(1)
+ val resultBucket =
handle.outboundThrottleMode.get().asInstanceOf[TokenBucket]
+ resultBucket.availableTokens should ===(50)
+ resultBucket.capacity should ===(100)
+ resultBucket.tokensPerSecond should ===(0.0)
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
