feat: honor LightDataFetcher on by level batching logic (#1795)

### 📝 Description
Now that
[PropertyDataFetcher](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/generator/graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/generator/execution/PropertyDataFetcher.kt#L31)
implements `LightDataFetcher` we need to make sure we respect that and
on the `LEVEL_DISPATCHED` batching mechanism, decorate the DataFetcher
properly depending if its `LightDataFetcher`, or just `DataFetcher`

Author: samuelAndalon

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/ExecutionBatchState.kt

@@ -17,6 +17,7 @@
 package com.expediagroup.graphql.dataloader.instrumentation.level.state
 
 import graphql.schema.DataFetcher
+import graphql.schema.LightDataFetcher
 
 enum class LevelState { NOT_DISPATCHED, DISPATCHED }
 
@@ -50,7 +51,7 @@ class ExecutionBatchState(documentHeight: Int) {
         *Array(documentHeight) { number -> Level(number + 1) to 0 }
     )
 
-    private val manuallyCompletableDataFetchers: MutableMap<Level, MutableList<ManuallyCompletableDataFetcher>> =
+    private val manuallyCompletableDataFetchers: MutableMap<Level, MutableList<ManualDataFetcher>> =
         mutableMapOf(
             *Array(documentHeight) { number -> Level(number + 1) to mutableListOf() }
         )
@@ -59,7 +60,7 @@ class ExecutionBatchState(documentHeight: Int) {
      * Check if the [ExecutionBatchState] contains a level
      *
      * @param level to check if his state is being calculated
-     * @return whether or not state contains the level
+     * @return whether state contains the level
      */
     fun contains(level: Level): Boolean = levelsState.containsKey(level)
 
@@ -117,16 +118,22 @@ class ExecutionBatchState(documentHeight: Int) {
      * @param dataFetcher to be instrumented
      * @return instrumented dataFetcher
      */
-    fun toManuallyCompletableDataFetcher(level: Level, dataFetcher: DataFetcher<*>): ManuallyCompletableDataFetcher =
-        ManuallyCompletableDataFetcher(dataFetcher).also { manuallyCompletableDataFetchers[level]?.add(it) }
+    fun toManuallyCompletableDataFetcher(level: Level, dataFetcher: DataFetcher<*>): ManualDataFetcher {
+        val manualDataFetcher = when (dataFetcher) {
+            is LightDataFetcher<*> -> ManuallyCompletableLightDataFetcher(dataFetcher)
+            else -> ManuallyCompletableDataFetcher(dataFetcher)
+        }
+        manuallyCompletableDataFetchers[level]?.add(manualDataFetcher)
+        return manualDataFetcher
+    }
 
     /**
      * Complete all the [manuallyCompletableDataFetchers]
      *
      * @param level which level should complete dataFetchers
      */
     fun completeDataFetchers(level: Level) {
-        manuallyCompletableDataFetchers[level]?.forEach(ManuallyCompletableDataFetcher::complete)
+        manuallyCompletableDataFetchers[level]?.forEach(ManualDataFetcher::complete)
     }
 
     /**

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/ManualDataFetcher.kt

@@ -0,0 +1,44 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.expediagroup.graphql.dataloader.instrumentation.level.state
+
+import graphql.schema.DataFetcher
+import java.util.concurrent.CompletableFuture
+
+/**
+ * DataFetcher Decorator that allows manual completion of dataFetchers
+ */
+abstract class ManualDataFetcher : DataFetcher<CompletableFuture<Any?>> {
+    val manualFuture: CompletableFuture<Any?> = CompletableFuture()
+    var originalFuture: CompletableFuture<Any?>? = null
+    var originalExpressionException: Exception? = null
+
+    /**
+     * Manually complete the [manualFuture] by handling the [originalFuture]
+     */
+    fun complete() {
+        when {
+            originalExpressionException != null -> manualFuture.completeExceptionally(originalExpressionException)
+            else -> originalFuture?.handle { result, exception ->
+                when {
+                    exception != null -> manualFuture.completeExceptionally(exception)
+                    else -> manualFuture.complete(result)
+                }
+            }
+        }
+    }
+}

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/ManuallyCompletableDataFetcher.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.
@@ -31,16 +31,11 @@ import java.util.concurrent.CompletableFuture
  */
 class ManuallyCompletableDataFetcher(
     private val originalDataFetcher: DataFetcher<*>
-) : DataFetcher<CompletableFuture<Any?>> {
-
-    private val manualFuture: CompletableFuture<Any?> = CompletableFuture()
-    private var originalFuture: CompletableFuture<Any?>? = null
-    private var originalExpressionException: Exception? = null
-
+) : ManualDataFetcher() {
     /**
-     * when attempting to get the value from dataFetcher, store the execute the [originalDataFetcher]
+     * when attempting to get the value from dataFetcher, execute the [originalDataFetcher]
      * and store the resulting future [originalFuture] and a possible [originalExpressionException] if
-     * an exception was thrown during the expression
+     * a synchronous exception was thrown during the execution
      *
      * @param environment dataFetchingEnvironment with information about the field
      * @return an uncompleted manualFuture that can be completed at later time
@@ -54,19 +49,4 @@ class ManuallyCompletableDataFetcher(
         }
         return manualFuture
     }
-
-    /**
-     * Manually complete the [manualFuture] by handling the [originalFuture]
-     */
-    fun complete() {
-        when {
-            originalExpressionException != null -> manualFuture.completeExceptionally(originalExpressionException)
-            else -> originalFuture?.handle { result, exception ->
-                when {
-                    exception != null -> manualFuture.completeExceptionally(exception)
-                    else -> manualFuture.complete(result)
-                }
-            }
-        }
-    }
 }

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/state/ManuallyCompletableLightDataFetcher.kt

@@ -0,0 +1,68 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.expediagroup.graphql.dataloader.instrumentation.level.state
+
+import graphql.execution.Async
+import graphql.schema.DataFetchingEnvironment
+import graphql.schema.GraphQLFieldDefinition
+import graphql.schema.LightDataFetcher
+import java.util.concurrent.CompletableFuture
+import java.util.function.Supplier
+
+/**
+ * LightDataFetcher Decorator that stores the original dataFetcher result (it's always a completable future)
+ * it stores the [originalFuture] as property and returns an uncompleted [manualFuture]
+ * then at later point manually call [complete] to complete the [manualFuture] with the [originalFuture] result
+ * to let ExecutionStrategy handle all futures
+ *
+ * @param originalDataFetcher original dataFetcher to be decorated
+ */
+class ManuallyCompletableLightDataFetcher(
+    private val originalDataFetcher: LightDataFetcher<*>
+) : ManualDataFetcher(), LightDataFetcher<CompletableFuture<Any?>> {
+
+    override fun get(environment: DataFetchingEnvironment): CompletableFuture<Any?> =
+        get(environment.fieldDefinition, environment.getSource()) { environment }
+
+    /**
+     * when attempting to get the value from LightDataFetcher, execute the [originalDataFetcher]
+     * and store the resulting future [originalFuture] and a possible [originalExpressionException] if
+     * a synchronous exception was thrown during the execution
+     *
+     * @param fieldDefinition the graphql field definition
+     * @param sourceObject the source object to get a value from
+     * @param environmentSupplier a supplier of the [DataFetchingEnvironment] that creates it lazily
+     * @return an uncompleted manualFuture that can be completed at later time
+     */
+    override fun get(
+        fieldDefinition: GraphQLFieldDefinition,
+        sourceObject: Any?,
+        environmentSupplier: Supplier<DataFetchingEnvironment>
+    ): CompletableFuture<Any?> {
+        try {
+            val fetchedValueRaw = originalDataFetcher.get(
+                fieldDefinition,
+                sourceObject,
+                environmentSupplier
+            )
+            originalFuture = Async.toCompletableFuture(fetchedValueRaw)
+        } catch (e: Exception) {
+            originalExpressionException = e
+        }
+        return manualFuture
+    }
+}