finagle-core: Share RetryBudget within a single StackClient

Problem

It mostly affects clients built with MethodBuilder but the fix is also beneficial for vanilla stack clients
should be put more RetryBudget-dependent stack modules into them in the future.

We're very inconsistent when it comes to RetryBudgets in MethodBuilders. Here is the current
state of things:

 - If a user configures their own RetryBudget, MethodBuilder will re-use it for all its methods (and
   within every method between its filters, retries and backups).
 - If a user doesn't configure their own RetryBudget, MethoddBiulder will create a new RetryBudget
   for each method. There is a 50% chance that the underlying client will share its RetryBudget with
   the *first* method created with MethodBuilder.

Solution

Always (unless already configured by a user) pre-configure a fresh copy of a RetryBudget for every
new client.

Differential Revision: https://phabricator.twitter.biz/D979325

Author: vkostyukov

build.sbt

@@ -360,7 +360,8 @@ lazy val finagleIntegration = Project(
     finagleNetty4Http,
     finagleRedis % "test",
     finagleThrift,
-    finagleThriftMux % "test->compile;test->test"
+    finagleThriftMux % "test->compile;test->test",
+    finaglePostgresql % "test->compile;test->test"
   )
 
 lazy val finagleToggle = Project(

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

@@ -5,7 +5,6 @@ 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.Retries
 import com.twitter.finagle.service.TimeoutFilter
 import com.twitter.finagle.stats.LazyStatsReceiver
 import com.twitter.finagle.stats.StatsReceiver
@@ -56,6 +55,7 @@ object MethodBuilder {
    */
   def from[Req, Rep](dest: Name, stackClient: StackClient[Req, Rep]): MethodBuilder[Req, Rep] = {
     val stack = modifiedStack(stackClient.stack)
+
     new MethodBuilder(
       new MethodPool[Req, Rep](stackClient.withStack(stack), dest, param.Label.Default),
       dest,
@@ -397,10 +397,7 @@ final class MethodBuilder[Req, Rep] private[finagle] (
       methodPool,
       dest,
       stack,
-      // If the RetryBudget is not configured, BackupRequestFilter and RetryFilter will each
-      // get a new instance of the default budget. Since we want them to share the same
-      // client retry budget, insert the budget into the params.
-      params + params[Retries.Budget],
+      params,
       config.copy(backup = brfParam)
     )
 

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

@@ -491,7 +491,9 @@ object StackClient {
     Stack.Params.empty +
       Stats(ClientStatsReceiver) +
       LoadBalancerFactory.HostStats(LoadedHostStatsReceiver) +
-      MetricBuilders(Some(new CoreMetricsRegistry()))
+      MetricBuilders(Some(new CoreMetricsRegistry())) +
+      // RetryBudget needs to be shared across one physical client.
+      Retries.Budget.default
 
   /**
    * A set of StackTransformers for transforming client stacks.

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

@@ -1120,8 +1120,8 @@ class MethodBuilderTest
     assert(client.params[Retries.Budget].retryBudget eq retryBudget)
   }
 
-  test("idempotent adds retry budget to params") {
-    val params = Stack.Params.empty
+  test("shares RetryBudget between methods") {
+    val params = StackClient.defaultParams
 
     val svc: Service[Int, Int] = Service.const(Future.value(1))
 
@@ -1131,16 +1131,16 @@ class MethodBuilderTest
     val stackClient = TestStackClient(stack, params)
     val client = MethodBuilder.from("mb", stackClient)
 
-    // there's no budget to start with; if we get the budget it's a new default one each time
-    assert(client.params[Retries.Budget].retryBudget ne client.params[Retries.Budget].retryBudget)
+    val method1 =
+      client.idempotent(1.percent, sendInterrupts = true, ResponseClassifier.Default)
 
-    val idempotentClient =
+    val method2 =
       client.idempotent(1.percent, sendInterrupts = true, ResponseClassifier.Default)
 
-    // now that the budget has been added (a new default one), we get the same one each time
+    // each method should share the same retryBudget
     assert(
-      idempotentClient.params[Retries.Budget].retryBudget eq
-        idempotentClient.params[Retries.Budget].retryBudget
+      method1.params[Retries.Budget].retryBudget eq
+        method2.params[Retries.Budget].retryBudget
     )
   }
 

finagle-integration/src/test/scala/BUILD

@@ -22,6 +22,7 @@ junit_tests(
         "finagle/finagle-mysql/src/main/scala",
         "finagle/finagle-mysql/src/test/scala/com/twitter/finagle/mysql/harness",
         "finagle/finagle-netty4-http",
+        "finagle/finagle-postgresql/src/main/scala/com/twitter/finagle",
         "finagle/finagle-redis",
         "finagle/finagle-thrift",
         "finagle/finagle-thriftmux",

finagle-integration/src/test/scala/com/twitter/finagle/integration/StackClientTest.scala

@@ -1,7 +1,15 @@
 package com.twitter.finagle.integration
 
+import com.twitter.finagle.service.Retries
 import com.twitter.finagle.service.StatsFilter
-import com.twitter.finagle.{Http, Memcached, Mysql, Redis, Thrift, ThriftMux}
+import com.twitter.finagle.Http
+import com.twitter.finagle.Memcached
+import com.twitter.finagle.Mysql
+import com.twitter.finagle.PostgreSql
+import com.twitter.finagle.Redis
+import com.twitter.finagle.Thrift
+import com.twitter.finagle.ThriftMux
+import com.twitter.finagle.client.StackBasedClient
 import org.scalatest.funsuite.AnyFunSuite
 
 class StackClientTest extends AnyFunSuite {
@@ -44,4 +52,23 @@ class StackClientTest extends AnyFunSuite {
     assert(!c.stack.contains(role))
   }
 
+  def testThatRetryBudgetIsShared(newClient: () => StackBasedClient[_, _], protocol: String): Unit =
+    test(s"$protocol clients shipped with retry-budget included") {
+      // every new clients gets new retry budget
+      assert(
+        newClient().params[Retries.Budget].retryBudget ne newClient()
+          .params[Retries.Budget].retryBudget)
+
+      // single client shares the retry budget between it's modules
+      val client = newClient()
+      assert(client.params[Retries.Budget].retryBudget eq client.params[Retries.Budget].retryBudget)
+    }
+
+  testThatRetryBudgetIsShared(Thrift.client _, "Thrift")
+  testThatRetryBudgetIsShared(ThriftMux.client _, "ThriftMux")
+  testThatRetryBudgetIsShared(Http.client _, "Http")
+  testThatRetryBudgetIsShared(Memcached.client _, "Memcached")
+  testThatRetryBudgetIsShared(Redis.client _, "Redis")
+  testThatRetryBudgetIsShared(Mysql.client _, "Mysql")
+  testThatRetryBudgetIsShared(PostgreSql.client _, "PostgreSql")
 }

finagle-postgresql/src/main/scala/com/twitter/finagle/PostgreSql.scala

@@ -40,7 +40,7 @@ object PostgreSql {
         NackAdmissionFilter.module[Request, Response]
       )
 
-  val defaultParams: Stack.Params = StackClient.defaultParams +
+  def defaultParams: Stack.Params = StackClient.defaultParams +
     // Keep NackAdmissionFilter disabled by default for backwards compatibility.
     NackAdmissionFilter.Disabled