Improve field wrapper performance (#2051)

ghostdogpr · Dec 19, 2023 · ee2d9e9 · ee2d9e9
1 parent 2548822
commit ee2d9e9
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 57 deletions.
diff --git a/core/src/main/scala/caliban/wrappers/ApolloTracing.scala b/core/src/main/scala/caliban/wrappers/ApolloTracing.scala
@@ -1,17 +1,18 @@
 package caliban.wrappers
 
-import java.time.format.DateTimeFormatter
-import java.time.{ Instant, ZoneId }
-import java.util.concurrent.TimeUnit
 import caliban.ResponseValue.{ ListValue, ObjectValue }
 import caliban.Value.{ IntValue, StringValue }
+import caliban._
 import caliban.execution.{ ExecutionRequest, FieldInfo }
 import caliban.parsing.adt.Document
 import caliban.rendering.DocumentRenderer
-import caliban.wrappers.Wrapper.{ EffectfulWrapper, FieldWrapper, OverallWrapper, ParsingWrapper, ValidationWrapper }
-import caliban.{ CalibanError, GraphQLRequest, GraphQLResponse, PathValue, ResponseValue }
+import caliban.wrappers.Wrapper._
 import zio._
-import zio.query.ZQuery
+import zio.query.{ UQuery, ZQuery }
+
+import java.time.format.DateTimeFormatter
+import java.time.{ Instant, ZoneId }
+import java.util.concurrent.TimeUnit
 
 object ApolloTracing {
 
@@ -32,14 +33,13 @@ object ApolloTracing {
     EffectfulWrapper(
       ZIO
         .whenZIO(isEnabledRef.get)(
-          Ref
-            .make(Tracing())
-            .map(ref =>
-              apolloTracingOverall(ref) |+|
-                apolloTracingParsing(ref) |+|
-                apolloTracingValidation(ref) |+|
-                apolloTracingField(ref, !excludePureFields)
-            )
+          for {
+            ref   <- Ref.make(Tracing())
+            clock <- ZIO.clock
+          } yield apolloTracingOverall(clock, ref) |+|
+            apolloTracingParsing(clock, ref) |+|
+            apolloTracingValidation(clock, ref) |+|
+            apolloTracingField(ZQuery.fromZIO(clock.nanoTime), ref, !excludePureFields)
         )
         .someOrElse(Wrapper.empty)
     )
@@ -122,19 +122,19 @@ object ApolloTracing {
       )
   }
 
-  private def apolloTracingOverall(ref: Ref[Tracing]): OverallWrapper[Any] =
+  private def apolloTracingOverall(clock: Clock, ref: Ref[Tracing]): OverallWrapper[Any] =
     new OverallWrapper[Any] {
       def wrap[R1](
         process: GraphQLRequest => ZIO[R1, Nothing, GraphQLResponse[CalibanError]]
       ): GraphQLRequest => ZIO[R1, Nothing, GraphQLResponse[CalibanError]] =
         (request: GraphQLRequest) =>
           for {
-            nanoTime    <- Clock.nanoTime
-            currentTime <- Clock.currentTime(TimeUnit.MILLISECONDS)
+            nanoTime    <- clock.nanoTime
+            currentTime <- clock.currentTime(TimeUnit.MILLISECONDS)
             _           <- ref.update(_.copy(startTime = currentTime, startTimeMonotonic = nanoTime))
             result      <- process(request).timed.flatMap { case (duration, result) =>
                              for {
-                               endTime <- Clock.currentTime(TimeUnit.MILLISECONDS)
+                               endTime <- clock.currentTime(TimeUnit.MILLISECONDS)
                                _       <- ref.update(_.copy(duration = duration, endTime = endTime))
                                tracing <- ref.get
                              } yield result.copy(
@@ -149,14 +149,14 @@ object ApolloTracing {
           } yield result
     }
 
-  private def apolloTracingParsing(ref: Ref[Tracing]): ParsingWrapper[Any] =
+  private def apolloTracingParsing(clock: Clock, ref: Ref[Tracing]): ParsingWrapper[Any] =
     new ParsingWrapper[Any] {
       def wrap[R1](
         process: String => ZIO[R1, CalibanError.ParsingError, Document]
       ): String => ZIO[R1, CalibanError.ParsingError, Document] =
         (query: String) =>
           for {
-            start              <- Clock.nanoTime
+            start              <- clock.nanoTime
             resultWithDuration <- process(query).timed
             (duration, result)  = resultWithDuration
             _                  <- ref.update(state =>
@@ -167,14 +167,14 @@ object ApolloTracing {
           } yield result
     }
 
-  private def apolloTracingValidation(ref: Ref[Tracing]): ValidationWrapper[Any] =
+  private def apolloTracingValidation(clock: Clock, ref: Ref[Tracing]): ValidationWrapper[Any] =
     new ValidationWrapper[Any] {
       def wrap[R1](
         process: Document => ZIO[R1, CalibanError.ValidationError, ExecutionRequest]
       ): Document => ZIO[R1, CalibanError.ValidationError, ExecutionRequest] =
         (doc: Document) =>
           for {
-            start              <- Clock.nanoTime
+            start              <- clock.nanoTime
             resultWithDuration <- process(doc).timed
             (duration, result)  = resultWithDuration
             _                  <- ref.update(state =>
@@ -186,33 +186,36 @@ object ApolloTracing {
           } yield result
     }
 
-  private def apolloTracingField(ref: Ref[Tracing], wrapPureValues: Boolean): FieldWrapper[Any] =
+  private def apolloTracingField(
+    nanoTime: UQuery[Long],
+    ref: Ref[Tracing],
+    wrapPureValues: Boolean
+  ): FieldWrapper[Any] =
     new FieldWrapper[Any](wrapPureValues) {
       def wrap[R1](
         query: ZQuery[R1, CalibanError.ExecutionError, ResponseValue],
         fieldInfo: FieldInfo
       ): ZQuery[R1, CalibanError.ExecutionError, ResponseValue] =
         for {
-          summarized            <- query.summarized(Clock.nanoTime)((_, _))
-          ((start, end), result) = summarized
-          duration               = Duration.fromNanos(end - start)
-          _                     <- ZQuery.fromZIO(
-                                     ref
-                                       .update(state =>
-                                         state.copy(
-                                           execution = state.execution.copy(
-                                             resolvers = Resolver(
-                                               path = fieldInfo.path,
-                                               parentType = fieldInfo.details.parentType.fold("")(DocumentRenderer.renderTypeName),
-                                               fieldName = fieldInfo.name,
-                                               returnType = DocumentRenderer.renderTypeName(fieldInfo.details.fieldType),
-                                               startOffset = start - state.startTimeMonotonic,
-                                               duration = duration
-                                             ) :: state.execution.resolvers
-                                           )
-                                         )
-                                       )
-                                   )
+          start  <- nanoTime
+          result <- query
+          end    <- nanoTime
+          _      <- ZQuery.fromZIO(
+                      ref.update(state =>
+                        state.copy(
+                          execution = state.execution.copy(
+                            resolvers = Resolver(
+                              path = fieldInfo.path,
+                              parentType = fieldInfo.details.parentType.fold("")(DocumentRenderer.renderTypeName),
+                              fieldName = fieldInfo.name,
+                              returnType = DocumentRenderer.renderTypeName(fieldInfo.details.fieldType),
+                              startOffset = start - state.startTimeMonotonic,
+                              duration = Duration.fromNanos(end - start)
+                            ) :: state.execution.resolvers
+                          )
+                        )
+                      )
+                    )
         } yield result
     }
 

diff --git a/core/src/main/scala/caliban/wrappers/FieldMetrics.scala b/core/src/main/scala/caliban/wrappers/FieldMetrics.scala
@@ -7,7 +7,7 @@ import caliban.wrappers.Wrapper.OverallWrapper
 import zio._
 import zio.metrics.MetricKeyType.Histogram
 import zio.metrics.{ Metric, MetricKey, MetricLabel }
-import zio.query.ZQuery
+import zio.query.{ UQuery, ZQuery }
 
 import scala.jdk.CollectionConverters._
 
@@ -60,8 +60,10 @@ object FieldMetrics {
       for {
         timings  <- Ref.make(List.empty[Timing])
         failures <- Ref.make(List.empty[String])
+        clock    <- ZIO.clock
         metrics   = new Metrics(totalLabel, durationLabel, buckets, extraLabels)
-      } yield overallWrapper(timings, failures, metrics) |+| fieldWrapper(timings, failures)
+      } yield overallWrapper(timings, failures, metrics) |+|
+        fieldWrapper(ZQuery.fromZIO(clock.nanoTime), timings, failures)
     )
 
   private def overallWrapper(
@@ -114,7 +116,11 @@ object FieldMetrics {
     map.asScala.toMap
   }
 
-  private def fieldWrapper(timings: Ref[List[Timing]], failures: Ref[List[String]]): Wrapper.FieldWrapper[Any] =
+  private def fieldWrapper(
+    nanoTime: UQuery[Long],
+    timings: Ref[List[Timing]],
+    failures: Ref[List[String]]
+  ): Wrapper.FieldWrapper[Any] =
     new Wrapper.FieldWrapper[Any] {
       def wrap[R](
         query: ZQuery[R, CalibanError.ExecutionError, ResponseValue],
@@ -138,10 +144,10 @@ object FieldMetrics {
           ZQuery.fromZIO(failures.update(fieldName :: _) *> ZIO.fail(e))
 
         for {
-          summarized        <- query.summarized(Clock.nanoTime)((s, e) => e - s).catchAll(recordFailure)
-          (duration, result) = summarized
-          timing             = makeTiming(duration)
-          _                 <- ZQuery.fromZIO(timings.update(timing :: _))
+          start  <- nanoTime
+          result <- query.catchAll(recordFailure)
+          end    <- nanoTime
+          _      <- ZQuery.fromZIO(timings.update(makeTiming(end - start) :: _))
         } yield result
       }
     }

diff --git a/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala b/federation/src/main/scala/caliban/federation/tracing/ApolloFederatedTracing.scala
@@ -31,7 +31,9 @@ object ApolloFederatedTracing {
       for {
         tracing <- Ref.make(Tracing(NodeTrie.empty))
         enabled <- Ref.make(false)
-      } yield apolloTracingOverall(tracing, enabled) |+| apolloTracingField(tracing, enabled, !excludePureFields)
+        clock   <- ZIO.clock
+      } yield apolloTracingOverall(clock, tracing, enabled) |+|
+        apolloTracingField(clock, tracing, enabled, !excludePureFields)
     )
 
   private def toTimestamp(epochMilli: Long): Timestamp =
@@ -40,7 +42,7 @@ object ApolloFederatedTracing {
       (epochMilli % 1000).toInt * 1000000
     )
 
-  private def apolloTracingOverall(ref: Ref[Tracing], enabled: Ref[Boolean]): OverallWrapper[Any] =
+  private def apolloTracingOverall(clock: Clock, ref: Ref[Tracing], enabled: Ref[Boolean]): OverallWrapper[Any] =
     new OverallWrapper[Any] {
       def wrap[R1](
         process: GraphQLRequest => ZIO[R1, Nothing, GraphQLResponse[CalibanError]]
@@ -54,11 +56,11 @@ object ApolloFederatedTracing {
             )
           )(
             for {
-              startNano             <- Clock.nanoTime
+              startNano             <- clock.nanoTime
               _                     <- ref.update(_.copy(startTime = startNano))
-              response              <- process(request).summarized(Clock.currentTime(TimeUnit.MILLISECONDS))((_, _))
+              response              <- process(request).summarized(clock.currentTime(TimeUnit.MILLISECONDS))((_, _))
               ((start, end), result) = response
-              endNano               <- Clock.nanoTime
+              endNano               <- clock.nanoTime
               tracing               <- ref.get
             } yield {
               val root = Trace(
@@ -82,7 +84,12 @@ object ApolloFederatedTracing {
           )
     }
 
-  private def apolloTracingField(ref: Ref[Tracing], enabled: Ref[Boolean], wrapPureValues: Boolean): FieldWrapper[Any] =
+  private def apolloTracingField(
+    clock: Clock,
+    ref: Ref[Tracing],
+    enabled: Ref[Boolean],
+    wrapPureValues: Boolean
+  ): FieldWrapper[Any] =
     new FieldWrapper[Any](wrapPureValues) {
       def wrap[R1](
         query: ZQuery[R1, CalibanError.ExecutionError, ResponseValue],
@@ -93,7 +100,7 @@ object ApolloFederatedTracing {
           .flatMap(
             if (_)
               for {
-                response                          <- query.either.summarized(Clock.nanoTime)((_, _))
+                response                          <- query.either.summarized(clock.nanoTime)((_, _))
                 ((startTime, endTime), summarized) = response
                 id                                 = Node.Id.ResponseName(fieldInfo.name)
                 result                            <- ZQuery.fromZIO(