He-Pin commented on code in PR #3097: URL: https://github.com/apache/pekko/pull/3097#discussion_r3444386873
########## remote/src/test/scala/org/apache/pekko/remote/transport/ThrottlerHandleSpec.scala: ########## @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * license agreements; and to You under the Apache License, version 2.0: + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * This file is part of the Apache Pekko project, which was derived from Akka. + */ + +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.TokenBucket +import pekko.testkit.PekkoSpec +import pekko.util.ByteString + +@nowarn("msg=deprecated") +class ThrottlerHandleSpec extends PekkoSpec { + + class TestHandle(writeResult: Boolean) 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 + 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) + } Review Comment: Consider adding two more test cases to strengthen the regression coverage: 1. **Successful write still consumes tokens** — the success path changed from `Boolean` to `Option[(ThrottleMode, ThrottleMode)]`. A test asserting that tokens are consumed and NOT refunded on a successful write would guard against a future regression that accidentally always-refunds. 2. **`Unthrottled` mode is unaffected** — with `Unthrottled` as the default throttle mode, a test that writes with `Unthrottled` set and verifies `wrappedHandle.write()` is called (and returns the correct boolean) would cover the most common code path. Example sketch: ```scala "consume tokens on successful write" 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 { val wrappedHandle = new TestHandle(writeResult = true) val handle = ThrottlerHandle(wrappedHandle, system.deadLetters) handle.write(ByteString("data")) should ===(true) wrappedHandle.writeAttempts should ===(1) } ``` Non-blocking — the fix itself is correct. These tests would further harden the regression net. ########## remote/src/main/scala/org/apache/pekko/remote/transport/ThrottlerTransportAdapter.scala: ########## @@ -627,21 +627,41 @@ private[transport] final case class ThrottlerHandle(_wrappedHandle: AssociationH override def write(payload: ByteString): Boolean = { val tokens = payload.length - @tailrec def tryConsume(currentBucket: ThrottleMode): Boolean = { + @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)) true + if (outboundThrottleMode.compareAndSet(currentBucket, newBucket)) Some((currentBucket, newBucket)) else tryConsume(outboundThrottleMode.get()) - } else false + } else None + } + + @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 => Review Comment: Minor observation: `bucket.tokensPerSecond == previous.tokensPerSecond` uses exact `Double` equality. This is safe here because both values originate from the same `SetThrottle` configuration (same `TokenBucket` instance lineage), so no floating-point drift is possible between them. If a future change were to compute `tokensPerSecond` independently, consider switching to a tolerance-based comparison. The CAS retry loop and the three-way dispatch logic are sound. Nice fix for a long-standing FIXME. -- 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]
