From 90838f12cae5d84a80fecda53bb24c385bfc3751 Mon Sep 17 00:00:00 2001
From: kyri-petrou <67301607+kyri-petrou@users.noreply.github.com>
Date: Sun, 24 Mar 2024 14:15:43 +0800
Subject: [PATCH] Improve performance of `ApolloFederatedTracing` wrapper
 (#2167)

* Improve performance of federated tracing wrapper

* Make wrapper trace implicit

* PR comment

* Use cached `typeNameRepr`
---
 .../main/scala/caliban/wrappers/Wrapper.scala |  2 +-
 .../tracing/ApolloFederatedTracing.scala      | 67 +++++++++++--------
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/core/src/main/scala/caliban/wrappers/Wrapper.scala b/core/src/main/scala/caliban/wrappers/Wrapper.scala
index a6b2b5b10..551796c3f 100644
--- a/core/src/main/scala/caliban/wrappers/Wrapper.scala
+++ b/core/src/main/scala/caliban/wrappers/Wrapper.scala
@@ -39,7 +39,7 @@ sealed trait Wrapper[-R] extends GraphQLAspect[Nothing, R] { self =>
     that.withWrapper(self)
 
   // Disables tracing only for wrappers in the caliban package
-  final private[caliban] def trace: Trace = Trace.empty
+  final private[caliban] implicit def trace: Trace = Trace.empty
 }
 
 object Wrapper {
diff --git a/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala b/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala
index e3ee959ec..1f3cdce31 100644
--- a/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala
+++ b/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala
@@ -98,6 +98,33 @@ object ApolloFederatedTracing {
     wrapPureValues: Boolean
   ): FieldWrapper[Any] =
     new FieldWrapper[Any](wrapPureValues) {
+
+      private def updateState(startTime: Long, fieldInfo: FieldInfo, error: Option[ExecutionError]): Unit = {
+        val endTime = nanoTime
+        val path    = (PathValue.Key(fieldInfo.name) :: fieldInfo.path).toVector
+        val _       = ref.updateAndGet(state =>
+          state.copy(
+            root = state.root.insert(
+              path,
+              Node(
+                id = Node.Id.ResponseName(fieldInfo.name),
+                startTime = startTime - state.startTime,
+                endTime = endTime - state.startTime,
+                `type` = fieldInfo.details.fieldType.typeNameRepr,
+                parentType = fieldInfo.details.parentType.map(_.typeNameRepr) getOrElse "",
+                originalFieldName = fieldInfo.details.alias.map(_ => fieldInfo.details.name) getOrElse "",
+                error = error.map { e =>
+                  Error(
+                    e.getMessage(),
+                    location = e.locationInfo.map(l => Location(l.line, l.column)).toSeq
+                  )
+                }.toSeq
+              )
+            )
+          )
+        )
+      }
+
       def wrap[R1](
         query: ZQuery[R1, CalibanError.ExecutionError, ResponseValue],
         fieldInfo: FieldInfo
@@ -105,34 +132,18 @@ object ApolloFederatedTracing {
         if (enabled.get())
           ZQuery.suspend {
             val startTime = nanoTime
-            query.either.flatMap { result =>
-              ZQuery.fromEither {
-                val endTime = nanoTime
-                val path    = (PathValue.Key(fieldInfo.name) :: fieldInfo.path).toVector
-                val _       = ref.updateAndGet(state =>
-                  state.copy(
-                    root = state.root.insert(
-                      path,
-                      Node(
-                        id = Node.Id.ResponseName(fieldInfo.name),
-                        startTime = startTime - state.startTime,
-                        endTime = endTime - state.startTime,
-                        `type` = fieldInfo.details.fieldType.toType().toString,
-                        parentType = fieldInfo.details.parentType.map(_.toType().toString) getOrElse "",
-                        originalFieldName = fieldInfo.details.alias.map(_ => fieldInfo.details.name) getOrElse "",
-                        error = result.left.toOption.collectFirst { case e: ExecutionError =>
-                          Error(
-                            e.getMessage(),
-                            location = e.locationInfo.map(l => Location(l.line, l.column)).toSeq
-                          )
-                        }.toSeq
-                      )
-                    )
-                  )
-                )
-                result
-              }
-            }
+            query.foldQuery(
+              error =>
+                ZQuery.fail {
+                  updateState(startTime, fieldInfo, Some(error))
+                  error
+                },
+              value =>
+                ZQuery.succeed {
+                  updateState(startTime, fieldInfo, None)
+                  value
+                }
+            )
           }
         else query