finagle/finagle-core: Add retry context to retried requests

Problem

Currently, it's hard to attribute a surge in request volume and failures with a
retry storm -- it's a positive feedback loop. In a system with a complex call
chain, it can be difficult to know where to look for client-side retries, and these
still may not accurately reflect what the server ultimately sees due to multiplicative
retries at different levels.

Solution

Introduce a new context, `c.t.f.Retry`, to signal that a request is a retry. We can then
add server-side stats to be able to determine how many retries we receive (and look
at the against total request volume).

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

Author: jcrossley

CHANGELOG.rst

@@ -36,6 +36,7 @@ New Features
   context per request.
 * finagle-netty4: Added support for custom event loop implementations. ``PHAB_ID=D1185136``
 * finagle-core: Add low priority offload executor. ``PHAB_ID=D1189064``
+* finagle-core: A new context, `c.t.f.Retry` is set for a request if it is a retry. ``PHAB_ID=D1228127``
 
 
 Bug Fixes

finagle-base-http/src/main/scala/com/twitter/finagle/http/codec/context/HttpContext.scala

@@ -52,7 +52,8 @@ object HttpContext {
     Array[HttpContext](
       HttpDeadline,
       HttpRequeues,
-      HttpBackupRequest
+      HttpBackupRequest,
+      HttpRetry
     )
   }
 

finagle-base-http/src/main/scala/com/twitter/finagle/http/codec/context/HttpRetriesContext.scala renamed to finagle-base-http/src/main/scala/com/twitter/finagle/http/codec/context/HttpRequeues.scala

@@ -0,0 +1,20 @@
+package com.twitter.finagle.http.codec.context
+
+import com.twitter.finagle.context.Contexts
+import com.twitter.finagle.context.RetryContext
+import com.twitter.finagle.context.RetryContext.Retry
+import com.twitter.util.Try
+
+private object HttpRetry extends HttpContext {
+
+  type ContextKeyType = Retry
+  val key: Contexts.broadcast.Key[Retry] = RetryContext.Ctx
+
+  def toHeader(retry: Retry): String = {
+    "1"
+  }
+
+  def fromHeader(header: String): Try[Retry] = {
+    RetryContext.returnRetry
+  }
+}

finagle-base-http/src/main/scala/com/twitter/finagle/http/codec/context/HttpRetry.scala

@@ -5,6 +5,7 @@ import com.twitter.finagle.context.BackupRequest
 import com.twitter.finagle.context.Contexts
 import com.twitter.finagle.context.Deadline
 import com.twitter.finagle.context.Requeues
+import com.twitter.finagle.context.RetryContext
 import com.twitter.finagle.http.Message
 import com.twitter.finagle.http.Method
 import com.twitter.finagle.http.Request
@@ -157,6 +158,23 @@ class HttpContextTest extends AnyFunSuite {
     }
   }
 
+  test("Retry written matches read") {
+    val m = newMsg()
+    assert(!RetryContext.isRetry)
+    RetryContext.withRetry {
+      HttpContext.write(m)
+      assert(RetryContext.isRetry)
+
+      Contexts.broadcast.letClearAll {
+        assert(!RetryContext.isRetry)
+
+        HttpContext.read(m) {
+          assert(RetryContext.isRetry)
+        }
+      }
+    }
+  }
+
   test("invalid context header value causes context to not be set") {
     val m = newMsg()
     m.headerMap.set("Finagle-Ctx-com.twitter.finagle.foo", ",,,")

finagle-base-http/src/test/scala/com/twitter/finagle/http/codec/context/HttpContextTest.scala

@@ -0,0 +1,41 @@
+package com.twitter.finagle.context
+
+import com.twitter.io.Buf
+import com.twitter.util.Return
+import com.twitter.util.Try
+
+/**
+ * Value-less context indicating that this request is a Retry. We do not include the retry # because
+ * it can be misleading. Suppose that we have services A -> B -> C. A makes a request to B, which makes
+ * a request to C, which fails and is retried 3 times. All 3 retries fail, so A retries its request
+ * to B. When C receives this request, it's really the 4th retry of the request, but because we cannot
+ * pass context information back up the stack, we are unable to correctly set the # retries in the
+ * context -- to A, it's the first retry, and to B, it's the first attempt.
+ */
+private[finagle] object RetryContext {
+
+  final class Retry
+
+  private val retry = new Retry
+  val returnRetry = Return(retry)
+
+  private final class Context extends Contexts.broadcast.Key[Retry]("c.t.f.Retry") {
+    def marshal(value: Retry): Buf = Buf.Empty
+
+    def tryUnmarshal(buf: Buf): Try[Retry] = {
+      returnRetry
+    }
+  }
+
+  val Ctx: Contexts.broadcast.Key[Retry] = new Context
+
+  def isRetry: Boolean = {
+    Contexts.broadcast.get(Ctx).nonEmpty
+  }
+
+  def withRetry[T](f: => T): T = {
+    Contexts.broadcast.let(Ctx, retry) {
+      f
+    }
+  }
+}

finagle-core/src/main/scala/com/twitter/finagle/context/RetryContext.scala

@@ -2,10 +2,17 @@ package com.twitter.finagle.service
 
 import com.twitter.conversions.DurationOps._
 import com.twitter.finagle.Filter.TypeAgnostic
+import com.twitter.finagle.context
 import com.twitter.finagle.param.HighResTimer
-import com.twitter.finagle.stats.{NullStatsReceiver, StatsReceiver}
-import com.twitter.finagle.tracing.{Annotation, Trace, TraceId}
-import com.twitter.finagle.{Backoff, FailureFlags, Filter, Service}
+import com.twitter.finagle.stats.NullStatsReceiver
+import com.twitter.finagle.stats.StatsReceiver
+import com.twitter.finagle.tracing.Annotation
+import com.twitter.finagle.tracing.Trace
+import com.twitter.finagle.tracing.TraceId
+import com.twitter.finagle.Backoff
+import com.twitter.finagle.FailureFlags
+import com.twitter.finagle.Filter
+import com.twitter.finagle.Service
 import com.twitter.util._
 
 object RetryingService {
@@ -101,7 +108,14 @@ class RetryFilter[Req, Rep](
       trace.recordRpc("retry")
     }
 
-    val svcRep = service(req)
+    val svcRep = if (isRetriedRequest) {
+      context.RetryContext.withRetry {
+        service(req)
+      }
+    } else {
+      service(req)
+    }
+
     if (trace.isActivelyTracing) {
       // we always trace the exception
       svcRep.respond {

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

@@ -0,0 +1,13 @@
+package com.twitter.finagle.context
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class RetryContextTest extends AnyFunSuite {
+
+  test("Get/Set") {
+    assert(!RetryContext.isRetry)
+    RetryContext.withRetry {
+      assert(RetryContext.isRetry)
+    }
+  }
+}

finagle-core/src/test/scala/com/twitter/finagle/context/RetryContextTest.scala

@@ -8,6 +8,7 @@ import com.twitter.finagle.Failure
 import com.twitter.finagle.FailureFlags
 import com.twitter.finagle.Service
 import com.twitter.finagle.WriteException
+import com.twitter.finagle.context.RetryContext
 import com.twitter.util._
 import org.mockito.ArgumentMatchers.anyObject
 import org.mockito.Mockito.times
@@ -92,6 +93,40 @@ class RetryFilterTest extends AnyFunSpec with MockitoSugar with BeforeAndAfter {
 
   describe("RetryFilter") {
 
+    it("sets retry context correctly in RetryFilter") {
+      val timer = new MockTimer
+
+      val stats = new InMemoryStatsReceiver()
+      val service = mock[Service[Int, Int]]
+      when(service.close(anyObject[Time])) thenReturn Future.Done
+      when(service(123))
+        .thenAnswer { _ =>
+          if (RetryContext.isRetry) Future(321)
+          else Future.exception(WriteException(new Exception("first attempt")))
+        }
+
+      val policy = RetryPolicy.tries[Try[Nothing]](
+        2,
+        {
+          case Throw(WriteException(_)) => true
+        })
+      val filter = new RetryExceptionsFilter[Int, Int](policy, timer, stats)
+      val retryingService = filter andThen service
+
+      Time.withCurrentTimeFrozen { tc =>
+        val result = retryingService(123)
+        verify(service, times(1))(123)
+        assert(!result.isDefined)
+
+        tc.advance(1.second)
+        timer.tick()
+
+        verify(service, times(2))(123)
+        assert(Await.result(result, 5.seconds) == 321)
+        assert(stats.stat("retries")() == Seq(1))
+      }
+    }
+
     it("respects RetryBudget") {
       val stats = new InMemoryStatsReceiver()