From c838d2bdaf6ab9d54851efe3aeaa31ea638d2b07 Mon Sep 17 00:00:00 2001
From: Jillian Crossley <jillianc@twitter.com>
Date: Wed, 13 Nov 2024 15:19:37 +0000
Subject: [PATCH] finagle/finagle-core: VerboseRequestTracer filter that adds
 tracing through the finagle stack

Problem

We want to be able to trace requests throughout the finagle stack during
testing.

Solution

Add a filter, VerboseRequestTracer, that adds this functionality for traced
requests. It is disabled by default and can be enabled via flag
`com.twitter.finagle.filter.verboseRequestTracing`.

Differential Revision: https://phabricator.twitter.biz/D1182361
---
 .../client/EndpointerStackClient.scala        |  9 ++-
 .../finagle/filter/VerboseRequestTracer.scala | 77 +++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 finagle-core/src/main/scala/com/twitter/finagle/filter/VerboseRequestTracer.scala

diff --git a/finagle-core/src/main/scala/com/twitter/finagle/client/EndpointerStackClient.scala b/finagle-core/src/main/scala/com/twitter/finagle/client/EndpointerStackClient.scala
index 46af55b7818..2863ab5c9d7 100644
--- a/finagle-core/src/main/scala/com/twitter/finagle/client/EndpointerStackClient.scala
+++ b/finagle-core/src/main/scala/com/twitter/finagle/client/EndpointerStackClient.scala
@@ -3,6 +3,7 @@ package com.twitter.finagle.client
 import com.twitter.finagle._
 import com.twitter.finagle.client.EndpointerStackClient.DimensionalClientScopes
 import com.twitter.finagle.filter.RequestLogger
+import com.twitter.finagle.filter.VerboseRequestTracer
 import com.twitter.finagle.naming.BindingFactory
 import com.twitter.finagle.param._
 import com.twitter.finagle.stack.nilStack
@@ -135,13 +136,19 @@ trait EndpointerStackClient[Req, Rep, This <: EndpointerStackClient[Req, Rep, Th
 
     val originalStack = {
       val baseStack = stack ++ (endpointer +: nilStack)
-      params[RequestLogger.Param] match {
+      val stackWithRequestTracing = params[RequestLogger.Param] match {
         case RequestLogger.Param.Enabled =>
           val transformer = RequestLogger.newStackTransformer(clientLabel)
           transformer(baseStack)
         case RequestLogger.Param.Disabled =>
           baseStack
       }
+      params[VerboseRequestTracer.Param] match {
+        case VerboseRequestTracer.Param.Enabled =>
+          VerboseRequestTracer.stackTransformer(stackWithRequestTracing)
+        case VerboseRequestTracer.Param.Disabled =>
+          stackWithRequestTracing
+      }
     }
 
     val transformedStack =
diff --git a/finagle-core/src/main/scala/com/twitter/finagle/filter/VerboseRequestTracer.scala b/finagle-core/src/main/scala/com/twitter/finagle/filter/VerboseRequestTracer.scala
new file mode 100644
index 00000000000..95b5561931c
--- /dev/null
+++ b/finagle-core/src/main/scala/com/twitter/finagle/filter/VerboseRequestTracer.scala
@@ -0,0 +1,77 @@
+package com.twitter.finagle.filter
+
+import com.twitter.app.GlobalFlag
+import com.twitter.finagle.ClientConnection
+import com.twitter.finagle.Service
+import com.twitter.finagle.ServiceFactory
+import com.twitter.finagle.ServiceFactoryProxy
+import com.twitter.finagle.ServiceProxy
+import com.twitter.finagle.Stack
+import com.twitter.finagle.tracing.Trace
+import com.twitter.util.Future
+
+object verboseRequestTracing
+    extends GlobalFlag[Boolean](
+      """Experimental flag. Enables verbose request tracing, which includes tracing though the finagle stack""".stripMargin
+    )
+
+private[twitter] object VerboseRequestTracer {
+
+  sealed trait Param {
+    def mk(): (Param, Stack.Param[Param]) = (this, Param.param)
+  }
+
+  private[finagle] object Param {
+    case object Disabled extends Param
+    case object Enabled extends Param
+
+    implicit val param: Stack.Param[Param] = new Stack.Param[Param] {
+      lazy val default: Param = {
+        verboseRequestTracing.get match {
+          case Some(value) if value => Enabled
+          case _ => Disabled
+        }
+      }
+    }
+  }
+
+  /**
+   * Enables the [[VerboseRequestTracer]].
+   */
+  val Enabled: Param = Param.Enabled
+
+  /**
+   * Disables the [[VerboseRequestTracer]] (disabled by default).
+   */
+  val Disabled: Param = Param.Disabled
+
+  private[finagle] val stackTransformer: Stack.Transformer =
+    new Stack.Transformer {
+      def apply[Req, Rep](stack: Stack[ServiceFactory[Req, Rep]]): Stack[ServiceFactory[Req, Rep]] =
+        stack.map((hd, sf) => withRequestTracing(hd.role, sf))
+    }
+
+  private[this] def withRequestTracing[Req, Rep](
+    role: Stack.Role,
+    svcFac: ServiceFactory[Req, Rep]
+  ): ServiceFactory[Req, Rep] =
+    new ServiceFactoryProxy[Req, Rep](svcFac) {
+      override def apply(conn: ClientConnection): Future[Service[Req, Rep]] = {
+        super.apply(conn).map { svc =>
+          new ServiceProxy[Req, Rep](svc) {
+            override def apply(request: Req): Future[Rep] = {
+              if (!Trace.isActivelyTracing) {
+                super.apply(request)
+              } else {
+                Trace.traceLocalFuture(role.name + "_async") {
+                  Trace.traceLocal(role.name + "_sync") {
+                    super.apply(request)
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+}