Preserve throttler tokens on failed writes

[fallintoplace]

Summary

ThrottlerHandle used to consume token-bucket capacity before delegating to the wrapped transport handle. If the wrapped handle returned false, no write happened but the bucket stayed depleted.

This restores or refunds the reserved token state when the wrapped write fails, while leaving the successful write and blackhole paths unchanged.

Tests

• remote/testOnly org.apache.pekko.remote.transport.ThrottlerHandleSpec org.apache.pekko.remote.classic.transport.ThrottleModeSpec org.apache.pekko.remote.classic.transport.ThrottlerTransportAdapterSpec
• remote/headerCheck
• remote/scalafmtAll

[He-Pin]

Use a dedicated class instead of tuple2

[He-Pin]

Review Summary

Verdict: Approve ✅

The fix correctly addresses the documented FIXME (#2825) where ThrottlerHandle.write consumed token bucket capacity even when the wrapped write returned false.

Correctness Analysis

1. tryConsume returning Option[(ThrottleMode, ThrottleMode)] — Sound. Returns the (previousBucket, consumedBucket) pair needed for a precise CAS-based refund when no concurrent mutation occurred.

2. refundTokens CAS loop — Correct three-way dispatch:

   • currentBucket == consumedBucket → full restore (no concurrent mutation)
   • Same TokenBucket config → add back tokens, cap at capacity
   • Mode changed (e.g. SetThrottle to different mode) → no-op fallback

3. @tailrec retry on CAS failure — Correct for lock-free concurrent safety.

4. Unthrottled / Blackhole paths — Blackhole short-circuits before consume (unchanged). Unthrottled flows through case bucket @ _ → tryConsume succeeds → refund's case _ falls through to no-op (nothing to refund). Both correct.

5. TOCTOU window between tryConsume and refundTokens — Accepted tradeoff: if throttle mode changes (via SetThrottle) during the window, the refund is lost. This is rare and benign.

Test Verification

• ThrottlerHandleSpec — 1 test, passes ✅
• ThrottleModeSpec — 5 tests, passes ✅
• ThrottlerTransportAdapterSpec — 2 tests, passes ✅

Suggestions (non-blocking)

See inline comments for test coverage improvements.

[He-Pin]

welcome and please fix license header

[He-Pin]

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:

"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.

[He-Pin]

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.

[pjfanning]

blocking until we have a CI run

[pjfanning]

build passed

unblocking