Improve performance of ApolloFederatedTracing wrapper (#2167)

* Improve performance of federated tracing wrapper * Make wrapper trace implicit * PR comment * Use cached `typeNameRepr`
ghostdogpr · Mar 24, 2024 · 90838f1 · 90838f1
1 parent ceaad5b
commit 90838f1
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 29 deletions.
diff --git 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
@@ -98,41 +98,52 @@ 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
       ): ZQuery[R1, CalibanError.ExecutionError, ResponseValue] =
         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