finagle-core: Deposit budget once in MethodBuilder

Problem

The methodbuilder interface applies the same retrybudget for RetryFilter and
RequeueFilter in the same stack. Resulting in double deposit of retrybudget.

Solution

Re-use the WithdrawOnlyRetryBudget to prevent depositing successes in the
RequeueFilter and only deposit in the RetryFilter.
Same solution as the `Retries.moduleWithRetryPolicy` takes.

Result

Retrybudget only gets deposited once for a request.

finagle-core/src/main/scala/com/twitter/finagle/client/MethodBuilder.scala

@@ -2,10 +2,7 @@ package com.twitter.finagle.client
 
 import com.twitter.finagle.Filter.TypeAgnostic
 import com.twitter.finagle.client.MethodBuilderTimeout.TunableDuration
-import com.twitter.finagle.service.Filterable
-import com.twitter.finagle.service.ResponseClass
-import com.twitter.finagle.service.ResponseClassifier
-import com.twitter.finagle.service.TimeoutFilter
+import com.twitter.finagle.service.{Filterable, ResponseClass, ResponseClassifier, Retries, RetryBudget, TimeoutFilter}
 import com.twitter.finagle.stats.LazyStatsReceiver
 import com.twitter.finagle.stats.StatsReceiver
 import com.twitter.finagle.tracing.TraceInitializerFilter
@@ -16,8 +13,10 @@ import com.twitter.finagle.Name
 import com.twitter.finagle.Service
 import com.twitter.finagle.ServiceFactory
 import com.twitter.finagle.Stack
+import com.twitter.finagle.Stack.Param
 import com.twitter.finagle.param
 import com.twitter.finagle._
+import com.twitter.finagle.service.Retries.WithdrawOnlyRetryBudget
 import com.twitter.util.tunable.Tunable
 import com.twitter.util.Closable
 import com.twitter.util.CloseOnce
@@ -539,6 +538,14 @@ final class MethodBuilder[Req, Rep] private[finagle] (
   }
 
   private[this] def materialize(): Unit = {
+    val params = if (this.config.retry.responseClassifier eq MethodBuilderRetry.Disabled) {
+      this.params
+    } else {
+      // add a withdrawonlybudget in the stack to prevent double deposit when retries are enabled on the methodbuilder
+      val currentBudget = this.params[Retries.Budget]
+      val wrappedBudget: RetryBudget = new WithdrawOnlyRetryBudget(currentBudget.retryBudget)
+      this.params + Retries.Budget(wrappedBudget, currentBudget.requeueBackoffs)
+    }
     methodPool.materialize(params)
   }
 

finagle-core/src/main/scala/com/twitter/finagle/client/MethodBuilderRetry.scala

@@ -105,7 +105,7 @@ private[finagle] class MethodBuilderRetry[Req, Rep] private[client] (mb: MethodB
 private[client] object MethodBuilderRetry {
   private[this] val DefaultMaxRetries = 2
 
-  private val Disabled: ResponseClassifier =
+  private[client] val Disabled: ResponseClassifier =
     ResponseClassifier.named("Disabled")(PartialFunction.empty)
 
   private def shouldRetry[Req, Rep](

finagle-core/src/main/scala/com/twitter/finagle/service/Retries.scala

@@ -104,7 +104,7 @@ object Retries {
    * the calls to `RetryBudget.request()` to count. This allows for
    * swallowing the call to `request` in the second filter.
    */
-  private class WithdrawOnlyRetryBudget(underlying: RetryBudget) extends RetryBudget {
+  private[finagle] class WithdrawOnlyRetryBudget(underlying: RetryBudget) extends RetryBudget {
     def deposit(): Unit = ()
     def tryWithdraw(): Boolean = underlying.tryWithdraw()
     def balance: Long = underlying.balance

finagle-core/src/test/scala/com/twitter/finagle/client/MethodBuilderRetryTest.scala

@@ -2,12 +2,14 @@ package com.twitter.finagle.client
 
 import com.twitter.conversions.DurationOps._
 import com.twitter.finagle.client.MethodBuilderTest.TestStackClient
-import com.twitter.finagle.service.{ReqRep, ResponseClass, _}
+import com.twitter.finagle.service._
 import com.twitter.finagle.stats.{InMemoryStatsReceiver, StatsReceiver}
-import com.twitter.finagle.{Failure, FailureFlags, Service, ServiceFactory, Stack, param}
+import com.twitter.finagle._
 import com.twitter.util.{Await, Future, Throw}
 import org.scalatest.funsuite.AnyFunSuite
 
+import java.util.concurrent.atomic.AtomicInteger
+
 class MethodBuilderRetryTest extends AnyFunSuite {
 
   private[this] class RetrySvc {
@@ -162,4 +164,64 @@ class MethodBuilderRetryTest extends AnyFunSuite {
     assert(ex.isFlagged(FailureFlags.Retryable))
     assert(stats.stat(clientName, "client", "retries")() == Seq(0))
   }
+
+  test("retries and requeues do not double deposit RetryBudget") {
+    val stats = new InMemoryStatsReceiver()
+    val budget = Retries.Budget(RetryBudget(ttl = 10.seconds, 10, 0.1))
+
+    val svc = Service.mk[Int, Int] { _ =>
+        Future.value(1)
+    }
+
+    val stack = Retries
+      .moduleRequeueable[Int, Int]
+      .toStack(Stack.leaf(Stack.Role("test"), ServiceFactory.const(svc)))
+    val ps =
+      Stack.Params.empty +
+        param.Label(clientName) +
+        param.Stats(stats) +
+        budget
+    val stackClient = TestStackClient(stack, ps)
+
+    val methodBuilder = MethodBuilder.from("retry_it", stackClient)
+    val client = methodBuilder
+      .newService("client")
+
+    assert(budget.retryBudget.balance == 100)
+    (1 to 10).foreach(_ => Await.result(client(1), 5.seconds))
+    assert(budget.retryBudget.balance == 101)
+
+    assert(stats.stat(clientName, "client", "retries")() == Seq.fill(10)(0))
+  }
+
+
+  test("MethodBuilder withRetry.disabled should allow RequeueFilter to deposit RetryBudget") {
+    val stats = new InMemoryStatsReceiver()
+    val budget = Retries.Budget(RetryBudget(ttl = 10.seconds, 10, 0.1))
+
+    val svc = Service.mk[Int, Int] { _ =>
+      Future.value(1)
+    }
+
+    val stack = Retries
+      .moduleRequeueable[Int, Int]
+      .toStack(Stack.leaf(Stack.Role("test"), ServiceFactory.const(svc)))
+    val ps =
+      Stack.Params.empty +
+        param.Label(clientName) +
+        param.Stats(stats) +
+        budget
+    val stackClient = TestStackClient(stack, ps)
+
+    val methodBuilder = MethodBuilder.from("retry_it", stackClient)
+    val client = methodBuilder
+      .withRetry.disabled
+      .newService("client")
+
+    assert(budget.retryBudget.balance == 100)
+    (1 to 10).foreach(_ => Await.result(client(1), 5.seconds))
+    assert(budget.retryBudget.balance == 101)
+
+    assert(stats.stat(clientName, "client", "retries")() == Seq.fill(10)(0))
+  }
 }