More timings

lib/src/main/java/graphql/nadel/engine/instrumentation/NadelInstrumentationTimer.kt

@@ -19,6 +19,10 @@ internal class NadelInstrumentationTimer(
         step: Step,
         function: () -> T,
     ): T {
+        if (!instrumentation.isTimingEnabled()) {
+            return function()
+        }
+
         val start = ticker()
 
         val result = try {
@@ -89,6 +93,10 @@ internal class NadelInstrumentationTimer(
         private var exception: Throwable? = null
 
         inline fun <T> time(step: Step, function: () -> T): T {
+            if (!instrumentation.isTimingEnabled()) {
+                return function()
+            }
+
             timings.computeIfAbsent(step) {
                 AtomicReference(Duration.ZERO)
             }

lib/src/main/java/graphql/nadel/engine/plan/NadelExecutionPlan.kt

@@ -2,6 +2,7 @@ package graphql.nadel.engine.plan
 
 import graphql.nadel.Service
 import graphql.nadel.engine.transform.NadelTransform
+import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters
 import graphql.normalized.ExecutableNormalizedField
 
 internal typealias AnyNadelExecutionPlanStep = NadelExecutionPlan.Step<Any>
@@ -14,6 +15,8 @@ data class NadelExecutionPlan(
         val service: Service,
         val field: ExecutableNormalizedField,
         val transform: NadelTransform<T>,
+        val queryTransformTimingStep: NadelInstrumentationTimingParameters.Step,
+        val resultTransformTimingStep: NadelInstrumentationTimingParameters.Step,
         val state: T,
     )
 

lib/src/main/java/graphql/nadel/engine/plan/NadelExecutionPlanFactory.kt

@@ -20,16 +20,32 @@ import graphql.nadel.engine.transform.skipInclude.NadelSkipIncludeTransform.Comp
 import graphql.nadel.hooks.NadelExecutionHooks
 import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.ChildStep
 import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.RootStep.ExecutionPlanning
+import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.RootStep.QueryTransforming
+import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters.RootStep.ResultTransforming
 import graphql.normalized.ExecutableNormalizedField
 
 internal class NadelExecutionPlanFactory(
     private val executionBlueprint: NadelOverallExecutionBlueprint,
-    private val transforms: List<NadelTransform<Any>>,
+    transforms: List<NadelTransform<Any>>,
 ) {
-    // This will avoid creating the ChildStep object too many times
+    /**
+     * This creates the [ChildStep] objects upfront to avoid constantly recreating them.
+     */
+    private data class TransformWithTimingInfo(
+        val transform: NadelTransform<Any>,
+        val executionPlanTimingStep: ChildStep,
+        val queryTransformTimingStep: ChildStep,
+        val resultTransformTimingStep: ChildStep,
+    )
+
     private val transformsWithTimingStepInfo = transforms
         .map { transform ->
-            transform to ChildStep(parent = ExecutionPlanning, transform = transform)
+            TransformWithTimingInfo(
+                transform = transform,
+                executionPlanTimingStep = ChildStep(parent = ExecutionPlanning, transform = transform),
+                queryTransformTimingStep = ChildStep(parent = QueryTransforming, transform = transform),
+                resultTransformTimingStep = ChildStep(parent = ResultTransforming, transform = transform),
+            )
         }
 
     /**
@@ -48,15 +64,16 @@ internal class NadelExecutionPlanFactory(
 
         executionContext.timer.batch { timer ->
             traverseQuery(rootField) { field ->
-                transformsWithTimingStepInfo.forEach { (transform, timingStep) ->
+                transformsWithTimingStepInfo.forEach { transformWithTimingInfo ->
+                    val transform = transformWithTimingInfo.transform
                     // This is a patch to prevent errors
                     // Ideally this should not happen but the proper fix requires more refactoring
                     // See NadelSkipIncludeTransform.isApplicable for more details
                     if (isSkipIncludeSpecialField(field) && ((transform as NadelTransform<*>) !is NadelSkipIncludeTransform)) {
                         return@forEach
                     }
 
-                    val state = timer.time(step = timingStep) {
+                    val state = timer.time(step = transformWithTimingInfo.executionPlanTimingStep) {
                         transform.isApplicable(
                             executionContext,
                             serviceExecutionContext,
@@ -71,10 +88,12 @@ internal class NadelExecutionPlanFactory(
                     if (state != null) {
                         executionSteps.add(
                             NadelExecutionPlan.Step(
-                                service,
-                                field,
-                                transform,
-                                state,
+                                service = service,
+                                field = field,
+                                transform = transform,
+                                queryTransformTimingStep = transformWithTimingInfo.queryTransformTimingStep,
+                                resultTransformTimingStep = transformWithTimingInfo.resultTransformTimingStep,
+                                state = state,
                             ),
                         )
                     }

lib/src/main/java/graphql/nadel/engine/transform/query/NadelQueryTransformer.kt

@@ -4,6 +4,7 @@ import graphql.nadel.Service
 import graphql.nadel.engine.NadelExecutionContext
 import graphql.nadel.engine.NadelServiceExecutionContext
 import graphql.nadel.engine.blueprint.NadelOverallExecutionBlueprint
+import graphql.nadel.engine.instrumentation.NadelInstrumentationTimer
 import graphql.nadel.engine.plan.NadelExecutionPlan
 import graphql.nadel.engine.transform.NadelTransform
 import graphql.nadel.engine.transform.NadelTransformFieldResult
@@ -17,6 +18,7 @@ class NadelQueryTransformer private constructor(
     private val serviceExecutionContext: NadelServiceExecutionContext,
     private val executionPlan: NadelExecutionPlan,
     private val transformContext: TransformContext,
+    private val timer: NadelInstrumentationTimer.BatchTimer,
 ) {
     companion object {
         suspend fun transformQuery(
@@ -29,24 +31,27 @@ class NadelQueryTransformer private constructor(
         ): TransformResult {
             val transformContext = TransformContext()
 
-            val transformer = NadelQueryTransformer(
-                executionBlueprint,
-                service,
-                executionContext,
-                serviceExecutionContext,
-                executionPlan,
-                transformContext,
-            )
-            val result = transformer.transform(field)
-                .also { rootFields ->
-                    transformer.fixParentRefs(parent = null, rootFields)
-                }
-
-            return TransformResult(
-                result = result,
-                artificialFields = transformContext.artificialFields,
-                overallToUnderlyingFields = transformContext.overallToUnderlyingFields,
-            )
+            executionContext.timer.batch().use { timer ->
+                val transformer = NadelQueryTransformer(
+                    executionBlueprint,
+                    service,
+                    executionContext,
+                    serviceExecutionContext,
+                    executionPlan,
+                    transformContext,
+                    timer,
+                )
+                val result = transformer.transform(field)
+                    .also { rootFields ->
+                        transformer.fixParentRefs(parent = null, rootFields)
+                    }
+
+                return TransformResult(
+                    result = result,
+                    artificialFields = transformContext.artificialFields,
+                    overallToUnderlyingFields = transformContext.overallToUnderlyingFields,
+                )
+            }
         }
     }
 
@@ -149,22 +154,24 @@ class NadelQueryTransformer private constructor(
     ): NadelTransformFieldResult {
         var fieldFromPreviousTransform: ExecutableNormalizedField = field
         var aggregatedTransformResult: NadelTransformFieldResult? = null
-        for ((_, _, transform, state) in transformationSteps) {
-            val transformResultForStep = transform.transformField(
-                executionContext,
-                serviceExecutionContext,
-                this,
-                executionBlueprint,
-                service,
-                fieldFromPreviousTransform,
-                state,
-            )
+        for (transformStep in transformationSteps) {
+            val transformResultForStep = timer.time(transformStep.queryTransformTimingStep) {
+                transformStep.transform.transformField(
+                    executionContext,
+                    serviceExecutionContext,
+                    this,
+                    executionBlueprint,
+                    service,
+                    fieldFromPreviousTransform,
+                    transformStep.state,
+                )
+            }
             aggregatedTransformResult = if (aggregatedTransformResult == null) {
                 transformResultForStep
             } else {
                 NadelTransformFieldResult(
-                    transformResultForStep.newField,
-                    aggregatedTransformResult.artificialFields + transformResultForStep.artificialFields,
+                    newField = transformResultForStep.newField,
+                    artificialFields = aggregatedTransformResult.artificialFields + transformResultForStep.artificialFields,
                 )
             }
             fieldFromPreviousTransform = transformResultForStep.newField ?: break

lib/src/main/java/graphql/nadel/engine/transform/result/NadelResultTransformer.kt

@@ -2,7 +2,6 @@ package graphql.nadel.engine.transform.result
 
 import graphql.GraphQLError
 import graphql.incremental.DeferPayload
-import graphql.nadel.NadelIncrementalServiceExecutionResult
 import graphql.nadel.Service
 import graphql.nadel.ServiceExecutionResult
 import graphql.nadel.engine.NadelExecutionContext
@@ -28,7 +27,7 @@ internal class NadelResultTransformer(private val executionBlueprint: NadelOvera
         artificialFields: List<ExecutableNormalizedField>,
         overallToUnderlyingFields: Map<ExecutableNormalizedField, List<ExecutableNormalizedField>>,
         service: Service,
-        result: ServiceExecutionResult
+        result: ServiceExecutionResult,
     ): ServiceExecutionResult {
         val nodes = JsonNodes(result.data)
         val instructions = getMutationInstructions(
@@ -53,7 +52,7 @@ internal class NadelResultTransformer(private val executionBlueprint: NadelOvera
         overallToUnderlyingFields: Map<ExecutableNormalizedField, List<ExecutableNormalizedField>>,
         service: Service,
         result: ServiceExecutionResult,
-        deferPayload: DeferPayload
+        deferPayload: DeferPayload,
     ): DeferPayload {
         val nodes = JsonNodes(
             deferPayload.getData<JsonMap?>() ?: emptyMap(),
@@ -81,31 +80,35 @@ internal class NadelResultTransformer(private val executionBlueprint: NadelOvera
         overallToUnderlyingFields: Map<ExecutableNormalizedField, List<ExecutableNormalizedField>>,
         service: Service,
         result: ServiceExecutionResult,
-        nodes: JsonNodes
+        nodes: JsonNodes,
     ): List<NadelResultInstruction> {
         val asyncInstructions = ArrayList<Deferred<List<NadelResultInstruction>>>()
 
         coroutineScope {
-            for ((field, steps) in executionPlan.transformationSteps) {
-                val underlyingFields = overallToUnderlyingFields[field]
-                if (underlyingFields.isNullOrEmpty()) continue
-
-                for (step in steps) {
-                    asyncInstructions.add(
-                        async {
-                            step.transform.getResultInstructions(
-                                executionContext,
-                                serviceExecutionContext,
-                                executionBlueprint,
-                                service,
-                                field,
-                                underlyingFields.first().parent,
-                                result,
-                                step.state,
-                                nodes
+            executionContext.timer.batch { timer ->
+                for ((field, steps) in executionPlan.transformationSteps) {
+                    val underlyingFields = overallToUnderlyingFields[field]
+                    if (underlyingFields.isNullOrEmpty()) continue
+
+                    for (step in steps) {
+                        timer.time(step.resultTransformTimingStep) {
+                            asyncInstructions.add(
+                                async {
+                                    step.transform.getResultInstructions(
+                                        executionContext,
+                                        serviceExecutionContext,
+                                        executionBlueprint,
+                                        service,
+                                        field,
+                                        underlyingFields.first().parent,
+                                        result,
+                                        step.state,
+                                        nodes
+                                    )
+                                }
                             )
                         }
-                    )
+                    }
                 }
             }
 
@@ -119,7 +122,6 @@ internal class NadelResultTransformer(private val executionBlueprint: NadelOvera
         return asyncInstructions.awaitAll().flatten()
     }
 
-
     private fun mutate(result: ServiceExecutionResult, instructions: List<NadelResultInstruction>) {
         instructions.forEach { transformation ->
             when (transformation) {

lib/src/main/java/graphql/nadel/instrumentation/ChainedNadelInstrumentation.kt

@@ -28,6 +28,12 @@ class ChainedNadelInstrumentation(
 ) : NadelInstrumentation {
     constructor(vararg instrumentations: NadelInstrumentation) : this(instrumentations.toList())
 
+    override fun isTimingEnabled(): Boolean {
+        return instrumentations.any {
+            it.isTimingEnabled()
+        }
+    }
+
     fun getInstrumentations(): List<NadelInstrumentation> {
         return Collections.unmodifiableList(instrumentations)
     }

lib/src/main/java/graphql/nadel/instrumentation/NadelInstrumentation.kt

@@ -28,6 +28,10 @@ import java.util.concurrent.CompletableFuture
  * operations such as fetching data and resolving it into objects.
  */
 interface NadelInstrumentation {
+    fun isTimingEnabled(): Boolean {
+        return false
+    }
+
     /**
      * This will be called just before execution to create an object that is given back to all instrumentation methods
      * to allow them to have per execution request state