feat(batching): avoid calculating document ast height in advance (#1800…

…) (#1801) ### 📝 Description cherry pick #1800
ExpediaGroup · Jun 19, 2023 · 5341525 · 5341525
1 parent 4ca522a
commit 5341525
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 73 deletions.
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/build.gradle.kts b/executions/graphql-kotlin-dataloader-instrumentation/build.gradle.kts
@@ -23,7 +23,7 @@ tasks {
                 limit {
                     counter = "INSTRUCTION"
                     value = "COVEREDRATIO"
-                    minimum = "0.93".toBigDecimal()
+                    minimum = "0.92".toBigDecimal()
                 }
                 limit {
                     counter = "BRANCH"

diff --git a/.../expediagroup/graphql/dataloader/instrumentation/extensions/ExecutionContextExtensions.kt b/.../expediagroup/graphql/dataloader/instrumentation/extensions/ExecutionContextExtensions.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,32 +16,8 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.extensions
 
-import graphql.analysis.QueryTraverser
-import graphql.analysis.QueryVisitorFieldEnvironment
 import graphql.execution.ExecutionContext
 import graphql.language.OperationDefinition
-import kotlin.math.max
-
-/**
- * Calculate the longest path of the [ExecutionContext] AST Document from the root node to a leaf node
- * @return the height of the AST Document
- */
-internal fun ExecutionContext.getDocumentHeight(): Int {
-    val getFieldDepth: (QueryVisitorFieldEnvironment?) -> Int = { queryVisitor ->
-        var hasQueryVisitor = queryVisitor
-        var height = 1
-        while (hasQueryVisitor != null) {
-            hasQueryVisitor = hasQueryVisitor.parentEnvironment
-            height++
-        }
-        height
-    }
-    return QueryTraverser.Builder().schema(graphQLSchema).document(document).variables(variables).build()
-        .reducePreOrder(
-            { queryVisitor, height -> max(getFieldDepth(queryVisitor.parentEnvironment), height) },
-            0
-        )
-}
 
 /**
  * Checks if the [ExecutionContext] is a [OperationDefinition.Operation.MUTATION]

diff --git a/...in/com/expediagroup/graphql/dataloader/instrumentation/level/state/ExecutionBatchState.kt b/...in/com/expediagroup/graphql/dataloader/instrumentation/level/state/ExecutionBatchState.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,37 +24,34 @@ enum class LevelState { NOT_DISPATCHED, DISPATCHED }
 /**
  * Handle the state of an [graphql.ExecutionInput]
  */
-class ExecutionBatchState(documentHeight: Int) {
-
-    private val levelsState: MutableMap<Level, LevelState> = mutableMapOf(
-        *Array(documentHeight) { number -> Level(number + 1) to LevelState.NOT_DISPATCHED }
-    )
-
-    private val expectedFetches: MutableMap<Level, Int> = mutableMapOf(
-        *Array(documentHeight) { number -> Level(number + 1) to 0 }
-    )
-    private val dispatchedFetches: MutableMap<Level, Int> = mutableMapOf(
-        *Array(documentHeight) { number -> Level(number + 1) to 0 }
-    )
-
-    private val expectedExecutionStrategies: MutableMap<Level, Int> = mutableMapOf(
-        *Array(documentHeight) { number ->
-            val level = Level(number + 1)
-            level to if (level.isFirst()) 1 else 0
+class ExecutionBatchState {
+    private val levelsState: MutableMap<Level, LevelState> = mutableMapOf()
+
+    private val expectedFetches: MutableMap<Level, Int> = mutableMapOf()
+    private val dispatchedFetches: MutableMap<Level, Int> = mutableMapOf()
+
+    private val expectedExecutionStrategies: MutableMap<Level, Int> = mutableMapOf()
+    private val dispatchedExecutionStrategies: MutableMap<Level, Int> = mutableMapOf()
+
+    private val onFieldValueInfos: MutableMap<Level, Int> = mutableMapOf()
+
+    private val manuallyCompletableDataFetchers: MutableMap<Level, MutableList<ManualDataFetcher>> = mutableMapOf()
+
+    /**
+     * Initializes a level state in the [ExecutionBatchState]
+     * @param level to be initialized
+     */
+    fun initializeLevelStateIfNeeded(level: Level) {
+        if (!this.contains(level)) {
+            levelsState[level] = LevelState.NOT_DISPATCHED
+            expectedFetches[level] = 0
+            dispatchedFetches[level] = 0
+            expectedExecutionStrategies[level] = if (level.isFirst()) 1 else 0
+            dispatchedExecutionStrategies[level] = 0
+            onFieldValueInfos[level] = 0
+            manuallyCompletableDataFetchers[level] = mutableListOf()
         }
-    )
-    private val dispatchedExecutionStrategies: MutableMap<Level, Int> = mutableMapOf(
-        *Array(documentHeight) { number -> Level(number + 1) to 0 }
-    )
-
-    private val onFieldValueInfos: MutableMap<Level, Int> = mutableMapOf(
-        *Array(documentHeight) { number -> Level(number + 1) to 0 }
-    )
-
-    private val manuallyCompletableDataFetchers: MutableMap<Level, MutableList<ManualDataFetcher>> =
-        mutableMapOf(
-            *Array(documentHeight) { number -> Level(number + 1) to mutableListOf() }
-        )
+    }
 
     /**
      * Check if the [ExecutionBatchState] contains a level

diff --git a/...ediagroup/graphql/dataloader/instrumentation/level/state/ExecutionLevelDispatchedState.kt b/...ediagroup/graphql/dataloader/instrumentation/level/state/ExecutionLevelDispatchedState.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.level.state
 
-import com.expediagroup.graphql.dataloader.instrumentation.extensions.getDocumentHeight
 import com.expediagroup.graphql.dataloader.instrumentation.extensions.getExpectedStrategyCalls
 import com.expediagroup.graphql.dataloader.instrumentation.level.execution.OnLevelDispatchedCallback
 import graphql.ExecutionInput
@@ -41,16 +40,16 @@ class ExecutionLevelDispatchedState(
     val executions = ConcurrentHashMap<ExecutionInput, ExecutionBatchState>()
 
     /**
-     * When a specific [ExecutionInput] starts his execution, calculate the height of the AST Document
+     * Initialize the [ExecutionBatchState] of this [ExecutionInput]
      *
      * @param parameters contains information of which [ExecutionInput] will start his execution
-     * @return a non null [InstrumentationContext] object
+     * @return a nullable [InstrumentationContext]
      */
     fun beginExecuteOperation(
         parameters: InstrumentationExecuteOperationParameters
     ): InstrumentationContext<ExecutionResult>? {
         executions.computeIfAbsent(parameters.executionContext.executionInput) {
-            ExecutionBatchState(parameters.executionContext.getDocumentHeight())
+            ExecutionBatchState()
         }
         return null
     }
@@ -71,6 +70,7 @@ class ExecutionLevelDispatchedState(
 
         executions.computeIfPresent(executionInput) { _, executionState ->
             executionState.also {
+                it.initializeLevelStateIfNeeded(level)
                 it.increaseExpectedFetches(level, fieldCount)
                 it.increaseDispatchedExecutionStrategies(level)
             }
@@ -174,18 +174,15 @@ class ExecutionLevelDispatchedState(
 
     /**
      * calculate if all executions sharing a graphQLContext was dispatched, by
-     * 1. Checking if the height of  all executions was already calculated.
-     * 2. Filter all executions sharing the same Level
-     * 3. check if all executions sharing the same level dispatched that level.
+     * 1. Checking if all executions started.
+     * 3. check if all executions dispatched the provided [level].
      *
      * @param level that execution state will be calculated
      * @return Boolean for allExecutionsDispatched statement
      */
     fun allExecutionsDispatched(level: Level): Boolean =
         executions
             .takeIf { executions -> executions.size == totalExecutions }
-            ?.filter { (_, executionState) -> executionState.contains(level) }
-            ?.takeIf { executionsWithSameLevel -> executionsWithSameLevel.isNotEmpty() }
             ?.all { (_, executionState) -> executionState.isLevelDispatched(level) }
             ?: false
 }
diff --git a/.../src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/Level.kt b/.../src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/Level.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -34,7 +34,7 @@ data class Level(private val number: Int) {
 
     /**
      * calculate if this [Level] is the first
-     * @return whether or not this is the first [Level]
+     * @return whether this is the first [Level]
      */
     fun isFirst(): Boolean = number == 1
 }
diff --git a/.../graphql/dataloader/instrumentation/level/DataLoaderLevelDispatchedInstrumentationTest.kt b/.../graphql/dataloader/instrumentation/level/DataLoaderLevelDispatchedInstrumentationTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2023 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -65,7 +65,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
         assertEquals(1, missionStatistics?.batchInvokeCount)
         assertEquals(2, missionStatistics?.batchLoadCount)
 
-        verify(exactly = 2) {
+        verify(exactly = 2 + ONE_LEVEL) {
             kotlinDataLoaderRegistry.dispatchAll()
         }
     }
@@ -96,7 +96,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
         assertEquals(1, missionStatistics?.batchInvokeCount)
         assertEquals(2, missionStatistics?.batchLoadCount)
 
-        verify(exactly = 3) {
+        verify(exactly = 3 + ONE_LEVEL) {
             kotlinDataLoaderRegistry.dispatchAll()
         }
     }
@@ -138,7 +138,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
         assertEquals(2, missionsByAstronautStatistics?.batchInvokeCount)
         assertEquals(2, missionsByAstronautStatistics?.batchLoadCount)
 
-        verify(exactly = 4) {
+        verify(exactly = 4 + ONE_LEVEL) {
             kotlinDataLoaderRegistry.dispatchAll()
         }
     }
@@ -160,4 +160,8 @@ class DataLoaderLevelDispatchedInstrumentationTest {
             dataLoaderLevelDispatchedInstrumentation.getOnLevelDispatchedCallback(ofType()) wasNot Called
         }
     }
+
+    companion object {
+        private const val ONE_LEVEL = 1
+    }
 }