Field wrapper optimizations (#2056)

* Optimize tracing field wrappers * Add implicit empty trace to Schema instead * Use `updateAndGet` on AtomicReference * Pre-determine string builder's size
ghostdogpr · Dec 24, 2023 · 01dba86 · 01dba86
1 parent 348ec3b
commit 01dba86
Show file tree

Hide file tree

Showing 10 changed files with 181 additions and 139 deletions.
diff --git a/core/src/main/scala/caliban/GraphQL.scala b/core/src/main/scala/caliban/GraphQL.scala
@@ -13,6 +13,7 @@ import caliban.schema._
 import caliban.validation.Validator
 import caliban.wrappers.Wrapper
 import caliban.wrappers.Wrapper._
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 import zio.{ IO, Trace, URIO, ZIO }
 
 /**
@@ -87,7 +88,9 @@ trait GraphQL[-R] { self =>
             _             <- Validator.validate(document, typeToValidate)
           } yield ()
 
-        private def checkHttpMethod(cfg: ExecutionConfiguration)(req: ExecutionRequest): IO[ValidationError, Unit] =
+        private def checkHttpMethod(cfg: ExecutionConfiguration)(req: ExecutionRequest)(implicit
+          trace: Trace
+        ): IO[ValidationError, Unit] =
           ZIO
             .when(req.operationType == OperationType.Mutation && !cfg.allowMutationsOverGetRequests) {
               HttpRequestMethod.get.flatMap {

diff --git a/core/src/main/scala/caliban/execution/Executor.scala b/core/src/main/scala/caliban/execution/Executor.scala
@@ -12,6 +12,7 @@ import caliban.schema.{ ReducedStep, Step, Types }
 import caliban.wrappers.Wrapper.FieldWrapper
 import zio._
 import zio.query._
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 import zio.stream.ZStream
 
 import scala.annotation.tailrec
@@ -210,9 +211,10 @@ object Executor {
           var remainingNames                         = names
           while (remainingResponses ne Nil) {
             val name = remainingNames.head
-            val resp = remainingResponses.head match {
-              case null => i -= 1; fromQueries(i)
-              case resp => resp
+            var resp = remainingResponses.head
+            if (resp eq null) {
+              i -= 1
+              resp = fromQueries(i)
             }
             results = (name, resp) :: results
             remainingResponses = remainingResponses.tail
@@ -338,7 +340,7 @@ object Executor {
     } yield response
   }
 
-  private[caliban] def fail(error: CalibanError): UIO[GraphQLResponse[CalibanError]] =
+  private[caliban] def fail(error: CalibanError)(implicit trace: Trace): UIO[GraphQLResponse[CalibanError]] =
     ZIO.succeed(GraphQLResponse(NullValue, List(error)))
 
   private[caliban] def mergeFields(field: Field, typeName: String): List[Field] = {

diff --git a/core/src/main/scala/caliban/introspection/adt/__Type.scala b/core/src/main/scala/caliban/introspection/adt/__Type.scala
@@ -5,6 +5,7 @@ import caliban.parsing.adt.Definition.TypeSystemDefinition.TypeDefinition
 import caliban.parsing.adt.Definition.TypeSystemDefinition.TypeDefinition._
 import caliban.parsing.adt.Type.{ ListType, NamedType }
 import caliban.parsing.adt.{ Directive, Type }
+import caliban.rendering.DocumentRenderer
 import caliban.schema.Annotations.GQLExcluded
 import caliban.schema.Types
 
@@ -24,6 +25,8 @@ case class __Type(
 ) { self =>
   final override lazy val hashCode: Int = super.hashCode()
 
+  private[caliban] lazy val typeNameRepr: String = DocumentRenderer.renderTypeName(this)
+
   def |+|(that: __Type): __Type = __Type(
     kind,
     (name ++ that.name).reduceOption((_, b) => b),

diff --git a/core/src/main/scala/caliban/parsing/Parser.scala b/core/src/main/scala/caliban/parsing/Parser.scala
@@ -4,7 +4,8 @@ import caliban.CalibanError.ParsingError
 import caliban.InputValue
 import caliban.parsing.adt._
 import fastparse._
-import zio.{ IO, ZIO }
+import zio.stacktracer.TracingImplicits.disableAutoTrace
+import zio.{ IO, Trace, ZIO }
 
 import scala.util.Try
 
@@ -14,7 +15,7 @@ object Parser {
   /**
    * Parses the given string into a [[caliban.parsing.adt.Document]] object or fails with a [[caliban.CalibanError.ParsingError]].
    */
-  def parseQuery(query: String): IO[ParsingError, Document] = {
+  def parseQuery(query: String)(implicit trace: Trace): IO[ParsingError, Document] = {
     val sm = SourceMapper(query)
     ZIO
       .attempt(parse(query, document(_)))

diff --git a/core/src/main/scala/caliban/schema/Schema.scala b/core/src/main/scala/caliban/schema/Schema.scala
@@ -12,8 +12,9 @@ import caliban.schema.Types._
 import caliban.uploads.Upload
 import caliban.{ InputValue, ResponseValue }
 import zio.query.ZQuery
+import zio.stacktracer.TracingImplicits.disableAutoTrace
 import zio.stream.ZStream
-import zio.{ Chunk, URIO, ZIO }
+import zio.{ Chunk, Trace, URIO, ZIO }
 
 import java.time._
 import java.time.format.DateTimeFormatter
@@ -81,6 +82,9 @@ trait Schema[-R, T] { self =>
    */
   def arguments: List[__InputValue] = Nil
 
+  // Disables traces for caliban-provided effectful schemas
+  private[caliban] implicit def trace: Trace = Trace.empty
+
   /**
    * Builds a new `Schema` of `A` from an existing `Schema` of `T` and a function from `A` to `T`.
    * @param f a function from `A` to `T`.
@@ -488,7 +492,7 @@ trait GenericSchema[R] extends SchemaDerivation[R] with TemporalSchema {
     }
 
   implicit def futureSchema[R0, A](implicit ev: Schema[R0, A]): Schema[R0, Future[A]]                                 =
-    effectSchema[R0, R0, R0, Throwable, A].contramap[Future[A]](future => ZIO.fromFuture(_ => future))
+    effectSchema[R0, R0, R0, Throwable, A].contramap[Future[A]](future => ZIO.fromFuture(_ => future)(Trace.empty))
   implicit def infallibleEffectSchema[R0, R1 >: R0, R2 >: R0, A](implicit ev: Schema[R2, A]): Schema[R0, URIO[R1, A]] =
     new Schema[R0, URIO[R1, A]] {
       override def optional: Boolean                                         = ev.optional

diff --git a/core/src/main/scala/caliban/validation/Validator.scala b/core/src/main/scala/caliban/validation/Validator.scala
@@ -21,7 +21,8 @@ import caliban.validation.Utils.isObjectType
 import caliban.{ Configurator, InputValue }
 import zio.prelude._
 import zio.prelude.fx.ZPure
-import zio.{ IO, ZIO }
+import zio.stacktracer.TracingImplicits.disableAutoTrace
+import zio.{ IO, Trace, ZIO }
 
 import scala.annotation.tailrec
 import scala.collection.mutable
@@ -47,13 +48,13 @@ object Validator {
   /**
    * Verifies that the given document is valid for this type. Fails with a [[caliban.CalibanError.ValidationError]] otherwise.
    */
-  def validate(document: Document, rootType: RootType): IO[ValidationError, Unit] =
+  def validate(document: Document, rootType: RootType)(implicit trace: Trace): IO[ValidationError, Unit] =
     Configurator.configuration.map(_.validations).flatMap(check(document, rootType, Map.empty, _).unit.toZIO)
 
   /**
    * Verifies that the given schema is valid. Fails with a [[caliban.CalibanError.ValidationError]] otherwise.
    */
-  def validateSchema[R](schema: RootSchemaBuilder[R]): IO[ValidationError, RootSchema[R]] = {
+  def validateSchema[R](schema: RootSchemaBuilder[R])(implicit trace: Trace): IO[ValidationError, RootSchema[R]] = {
     val schemaValidation = validateSchemaEither(schema)
     ZIO.fromEither(schemaValidation)
   }

diff --git a/core/src/main/scala/caliban/wrappers/ApolloTracing.scala b/core/src/main/scala/caliban/wrappers/ApolloTracing.scala
@@ -5,14 +5,14 @@ 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._
 import zio._
-import zio.query.{ UQuery, ZQuery }
+import zio.query.ZQuery
 
 import java.time.format.DateTimeFormatter
 import java.time.{ Instant, ZoneId }
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicReference
 
 object ApolloTracing {
 
@@ -34,12 +34,12 @@ object ApolloTracing {
       ZIO
         .whenZIO(isEnabledRef.get)(
           for {
-            ref   <- Ref.make(Tracing())
+            ref   <- ZIO.succeed(new AtomicReference(Tracing()))
             clock <- ZIO.clock
           } yield apolloTracingOverall(clock, ref) |+|
             apolloTracingParsing(clock, ref) |+|
             apolloTracingValidation(clock, ref) |+|
-            apolloTracingField(ZQuery.fromZIO(clock.nanoTime), ref, !excludePureFields)
+            Unsafe.unsafe(implicit u => apolloTracingField(clock.unsafe.nanoTime(), ref, !excludePureFields))
         )
         .someOrElse(Wrapper.empty)
     )
@@ -58,14 +58,14 @@ object ApolloTracing {
     .ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")
     .withZone(ZoneId.of("UTC"))
 
-  case class Parsing(startOffset: Long = 0, duration: Duration = Duration.Zero) {
+  case class Parsing(startOffset: Long = 0, durationNanos: Long = 0L) {
     def toResponseValue: ResponseValue =
-      ObjectValue(List("startOffset" -> IntValue(startOffset), "duration" -> IntValue(duration.toNanos)))
+      ObjectValue(List("startOffset" -> IntValue(startOffset), "duration" -> IntValue(durationNanos)))
   }
 
-  case class Validation(startOffset: Long = 0, duration: Duration = Duration.Zero) {
+  case class Validation(startOffset: Long = 0, durationNanos: Long = 0L) {
     def toResponseValue: ResponseValue =
-      ObjectValue(List("startOffset" -> IntValue(startOffset), "duration" -> IntValue(duration.toNanos)))
+      ObjectValue(List("startOffset" -> IntValue(startOffset), "duration" -> IntValue(durationNanos)))
   }
 
   case class Resolver(
@@ -74,7 +74,7 @@ object ApolloTracing {
     fieldName: String = "",
     returnType: String = "",
     startOffset: Long = 0,
-    duration: Duration = Duration.Zero
+    durationNanos: Long = 0
   ) {
     def toResponseValue: ResponseValue =
       ObjectValue(
@@ -84,26 +84,30 @@ object ApolloTracing {
           "fieldName"   -> StringValue(fieldName),
           "returnType"  -> StringValue(returnType),
           "startOffset" -> IntValue(startOffset),
-          "duration"    -> IntValue(duration.toNanos)
+          "duration"    -> IntValue(durationNanos)
         )
       )
   }
 
+  object Resolver {
+    implicit val ordering: Ordering[Resolver] = { (x: Resolver, y: Resolver) =>
+      val ord1 = Ordering.Long.compare(x.startOffset, y.startOffset)
+      if (ord1 != 0) ord1
+      else Ordering.Long.compare(x.durationNanos, y.durationNanos)
+    }
+  }
+
   case class Execution(resolvers: List[Resolver] = Nil) {
     def toResponseValue: ResponseValue =
-      ObjectValue(
-        List(
-          "resolvers" -> ListValue(resolvers.sortBy(r => (r.startOffset, r.duration.toNanos)).map(_.toResponseValue))
-        )
-      )
+      ObjectValue(List("resolvers" -> ListValue(resolvers.sorted.map(_.toResponseValue))))
   }
 
   case class Tracing(
     version: Int = 1,
     startTime: Long = 0,
     endTime: Long = 0,
     startTimeMonotonic: Long = 0,
-    duration: Duration = Duration.Zero,
+    durationNanos: Long = 0L,
     parsing: Parsing = Parsing(),
     validation: Validation = Validation(),
     execution: Execution = Execution()
@@ -114,15 +118,15 @@ object ApolloTracing {
           "version"    -> IntValue(version),
           "startTime"  -> StringValue(dateFormatter.format(Instant.ofEpochMilli(startTime))),
           "endTime"    -> StringValue(dateFormatter.format(Instant.ofEpochMilli(endTime))),
-          "duration"   -> IntValue(duration.toNanos),
+          "duration"   -> IntValue(durationNanos),
           "parsing"    -> parsing.toResponseValue,
           "validation" -> validation.toResponseValue,
           "execution"  -> execution.toResponseValue
         )
       )
   }
 
-  private def apolloTracingOverall(clock: Clock, ref: Ref[Tracing]): OverallWrapper[Any] =
+  private def apolloTracingOverall(clock: Clock, ref: AtomicReference[Tracing]): OverallWrapper[Any] =
     new OverallWrapper[Any] {
       def wrap[R1](
         process: GraphQLRequest => ZIO[R1, Nothing, GraphQLResponse[CalibanError]]
@@ -131,12 +135,11 @@ object ApolloTracing {
           for {
             nanoTime    <- clock.nanoTime
             currentTime <- clock.currentTime(TimeUnit.MILLISECONDS)
-            _           <- ref.update(_.copy(startTime = currentTime, startTimeMonotonic = nanoTime))
+            _           <- ZIO.succeed(ref.updateAndGet(_.copy(startTime = currentTime, startTimeMonotonic = nanoTime)))
             result      <- process(request).timed.flatMap { case (duration, result) =>
                              for {
                                endTime <- clock.currentTime(TimeUnit.MILLISECONDS)
-                               _       <- ref.update(_.copy(duration = duration, endTime = endTime))
-                               tracing <- ref.get
+                               tracing <- ZIO.succeed(ref.get.copy(durationNanos = duration.toNanos, endTime = endTime))
                              } yield result.copy(
                                extensions = Some(
                                  ObjectValue(
@@ -149,7 +152,7 @@ object ApolloTracing {
           } yield result
     }
 
-  private def apolloTracingParsing(clock: Clock, ref: Ref[Tracing]): ParsingWrapper[Any] =
+  private def apolloTracingParsing(clock: Clock, ref: AtomicReference[Tracing]): ParsingWrapper[Any] =
     new ParsingWrapper[Any] {
       def wrap[R1](
         process: String => ZIO[R1, CalibanError.ParsingError, Document]
@@ -159,15 +162,18 @@ object ApolloTracing {
             start              <- clock.nanoTime
             resultWithDuration <- process(query).timed
             (duration, result)  = resultWithDuration
-            _                  <- ref.update(state =>
-                                    state.copy(
-                                      parsing = state.parsing.copy(startOffset = start - state.startTimeMonotonic, duration = duration)
+            _                  <- ZIO.succeed(
+                                    ref.updateAndGet(state =>
+                                      state.copy(
+                                        parsing = state.parsing
+                                          .copy(startOffset = start - state.startTimeMonotonic, durationNanos = duration.toNanos)
+                                      )
                                     )
                                   )
           } yield result
     }
 
-  private def apolloTracingValidation(clock: Clock, ref: Ref[Tracing]): ValidationWrapper[Any] =
+  private def apolloTracingValidation(clock: Clock, ref: AtomicReference[Tracing]): ValidationWrapper[Any] =
     new ValidationWrapper[Any] {
       def wrap[R1](
         process: Document => ZIO[R1, CalibanError.ValidationError, ExecutionRequest]
@@ -177,46 +183,49 @@ object ApolloTracing {
             start              <- clock.nanoTime
             resultWithDuration <- process(doc).timed
             (duration, result)  = resultWithDuration
-            _                  <- ref.update(state =>
-                                    state.copy(
-                                      validation =
-                                        state.validation.copy(startOffset = start - state.startTimeMonotonic, duration = duration)
-                                    )
-                                  )
+            _                  <-
+              ZIO.succeed(
+                ref.updateAndGet(state =>
+                  state.copy(
+                    validation = state.validation
+                      .copy(startOffset = start - state.startTimeMonotonic, durationNanos = duration.toNanos)
+                  )
+                )
+              )
           } yield result
     }
 
   private def apolloTracingField(
-    nanoTime: UQuery[Long],
-    ref: Ref[Tracing],
+    nanoTime: => Long,
+    ref: AtomicReference[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 {
-          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
+        ZQuery.suspend {
+          val start = nanoTime
+          query.map { result =>
+            val end = nanoTime
+            val _   = ref.updateAndGet(state =>
+              state.copy(
+                execution = state.execution.copy(
+                  resolvers = Resolver(
+                    path = fieldInfo.path,
+                    parentType = fieldInfo.details.parentType.fold("")(_.typeNameRepr),
+                    fieldName = fieldInfo.name,
+                    returnType = fieldInfo.details.fieldType.typeNameRepr,
+                    startOffset = start - state.startTimeMonotonic,
+                    durationNanos = end - start
+                  ) :: state.execution.resolvers
+                )
+              )
+            )
+            result
+          }
+        }
     }
 
 }