feat: batching working again

Author: Samuel Vazquez

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/exceptions/MissingKotlinDataLoaderRegistryException.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2024 Expedia, Inc
+ * Copyright 2025 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,35 +16,31 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.extensions
 
-import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistry
 import com.expediagroup.graphql.dataloader.instrumentation.exceptions.MissingInstrumentationStateException
-import com.expediagroup.graphql.dataloader.instrumentation.exceptions.MissingKotlinDataLoaderRegistryException
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import graphql.schema.DataFetchingEnvironment
 import org.dataloader.DataLoader
+import org.dataloader.DataLoaderRegistry
 import java.util.concurrent.CompletableFuture
 
 /**
- * Check if all futures collected on [KotlinDataLoaderRegistry.dispatchAll] were handled
+ * Check if all futures collected on [DataLoaderRegistry.dispatchAll] were handled
  * and if we have more futures than we had when we started to dispatch, if so,
  * means that [DataLoader]s were chained, so we need to dispatch the dataLoaderRegistry.
+ *
+ * @throws MissingInstrumentationStateException if a [SyncExecutionExhaustedState] instance is not present in the graphQLContext
  */
 fun <V> CompletableFuture<V>.dispatchIfNeeded(
     environment: DataFetchingEnvironment
 ): CompletableFuture<V> {
-    val dataLoaderRegistry = environment.dataLoaderRegistry as? KotlinDataLoaderRegistry ?: throw MissingKotlinDataLoaderRegistryException()
+    val dataLoaderRegistry = environment.dataLoaderRegistry
+    val syncExecutionExhaustedState = environment
+        .graphQlContext
+        .get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
+        ?: throw MissingInstrumentationStateException()
 
-    if (dataLoaderRegistry.dataLoadersInvokedOnDispatch()) {
-        when {
-            environment.graphQlContext.hasKey(SyncExecutionExhaustedState::class) -> {
-                environment
-                    .graphQlContext.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
-                    .ifAllSyncExecutionsExhausted {
-                        dataLoaderRegistry.dispatchAll()
-                    }
-            }
-            else -> throw MissingInstrumentationStateException()
-        }
+    if (syncExecutionExhaustedState.dataLoadersInvokedAfterDispatch() && syncExecutionExhaustedState.allSyncExecutionsExhausted()) {
+        dataLoaderRegistry.dispatchAll()
     }
     return this
 }

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/extensions/CompletableFutureExtensions.kt

@@ -16,23 +16,29 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion
 
-import com.expediagroup.graphql.dataloader.KotlinDataLoaderInstrumentation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
+import org.dataloader.DataLoader
+import org.dataloader.instrumentation.DataLoaderInstrumentation
 import org.dataloader.instrumentation.DataLoaderInstrumentationContext
 
-internal class DataLoaderInstrumentationContext(
+class DataLoaderSyncExecutionExhaustedDataLoaderDispatcher(
     private val syncExecutionExhaustedState: SyncExecutionExhaustedState
-): DataLoaderInstrumentationContext<Any?> {
-    override fun onDispatched() {
-        syncExecutionExhaustedState.onDataLoaderPromiseDispatched()
-    }
-    override fun onCompleted(result: Any?, t: Throwable?) {
-        syncExecutionExhaustedState.onDataLoaderPromiseCompleted()
-    }
-}
+): DataLoaderInstrumentation {
+
+    private val contextForSyncExecutionExhausted: DataLoaderInstrumentationContext<Any?> =
+        object: DataLoaderInstrumentationContext<Any?> {
+            override fun onDispatched() {
+                syncExecutionExhaustedState.onDataLoaderPromiseDispatched()
+            }
+            override fun onCompleted(result: Any?, t: Throwable?) {
+                syncExecutionExhaustedState.onDataLoaderPromiseCompleted(result, t)
+            }
+        }
 
-class DataLoaderInstrumentationForSyncExecutionExhausted(
-    syncExecutionExhaustedState: SyncExecutionExhaustedState
-) : KotlinDataLoaderInstrumentation(
-    DataLoaderInstrumentationContext(syncExecutionExhaustedState)
-)
+    override fun beginLoad(
+        dataLoader: DataLoader<*, *>,
+        key: Any,
+        loadContext: Any?
+    ): DataLoaderInstrumentationContext<Any?> =
+        contextForSyncExecutionExhausted
+}

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderInstrumentationForSyncExecutionExhausted.kt renamed to executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedDataLoaderDispatcher.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2024 Expedia, Inc
+ * Copyright 2025 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,15 +14,13 @@
  * limitations under the License.
  */
 
-package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.execution
+package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion
 
 import com.expediagroup.graphql.dataloader.instrumentation.extensions.isMutation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
-import graphql.ExecutionInput
 import graphql.ExecutionResult
 import graphql.GraphQLContext
 import graphql.execution.ExecutionContext
-import graphql.execution.ExecutionId
 import graphql.execution.instrumentation.ExecuteObjectInstrumentationContext
 import graphql.execution.instrumentation.ExecutionStrategyInstrumentationContext
 import graphql.execution.instrumentation.FieldFetchingInstrumentationContext
@@ -32,41 +30,22 @@ import graphql.execution.instrumentation.InstrumentationState
 import graphql.execution.instrumentation.SimplePerformantInstrumentation
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
+import graphql.execution.instrumentation.parameters.InstrumentationFieldCompleteParameters
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
 
-/**
- * typealias that represents the signature of a callback that will be executed when sync execution is exhausted
- */
-internal typealias OnSyncExecutionExhaustedCallback = (List<ExecutionId>) -> Unit
-
 /**
  * Custom GraphQL [Instrumentation] that calculate the synchronous execution exhaustion
  * of all GraphQL operations sharing the same [GraphQLContext]
  */
-abstract class AbstractSyncExecutionExhaustedInstrumentation : SimplePerformantInstrumentation() {
-    /**
-     * This is invoked each time instrumentation attempts to calculate exhaustion state, this can be called from either
-     * `beginFieldField.dispatch` or `beginFieldFetch.complete`.
-     *
-     * @param parameters contains information of which [ExecutionInput] caused the calculation
-     * @return [OnSyncExecutionExhaustedCallback] to invoke when the synchronous execution of all operations was exhausted
-     */
-    abstract fun getOnSyncExecutionExhaustedCallback(
-        parameters: SyncExecutionExhaustedInstrumentationParameters
-    ): OnSyncExecutionExhaustedCallback
+class GraphQLSyncExecutionExhaustedDataLoaderDispatcher : SimplePerformantInstrumentation() {
 
     override fun beginExecution(
         parameters: InstrumentationExecutionParameters,
         state: InstrumentationState?
     ): InstrumentationContext<ExecutionResult>? =
         parameters.graphQLContext
             ?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
-            ?.beginExecution(
-                parameters,
-                this.getOnSyncExecutionExhaustedCallback(
-                    SyncExecutionExhaustedInstrumentationParameters(parameters.executionInput)
-                )
-            )
+            ?.beginExecution(parameters)
 
     override fun beginExecutionStrategy(
         parameters: InstrumentationExecutionStrategyParameters,
@@ -94,10 +73,13 @@ abstract class AbstractSyncExecutionExhaustedInstrumentation : SimplePerformantI
     ): FieldFetchingInstrumentationContext? =
         parameters.executionContext.takeUnless(ExecutionContext::isMutation)
             ?.graphQLContext?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
-            ?.beginFieldFetching(
-                parameters,
-                this.getOnSyncExecutionExhaustedCallback(
-                    SyncExecutionExhaustedInstrumentationParameters(parameters.executionContext.executionInput)
-                )
-            )
+            ?.beginFieldFetching(parameters)
+
+    override fun beginFieldCompletion(
+        parameters: InstrumentationFieldCompleteParameters,
+        state: InstrumentationState?
+    ): InstrumentationContext<Any>? {
+        println("field completed: ${parameters.fetchedValue}")
+        return super.beginFieldCompletion(parameters, state)
+    }
 }

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentation.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2025 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.syncexhaustion.state
+
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.atomic.AtomicInteger
+
+class DataLoaderDispatchState {
+
+    // count number of loads
+    private val onLoadCounter: AtomicInteger = AtomicInteger(0)
+    // take snapshot of load counts when dispatchAll is invoked, then on every load complete decrease it
+    private val onDispatchLoadCounter: AtomicInteger = AtomicInteger(0)
+
+    fun takeSnapshot() {
+        onDispatchLoadCounter.set(onLoadCounter.get())
+        onLoadCounter.set(0)
+    }
+
+    /**
+     * Will signal when all dependants of all [onDispatchFutures] were invoked,
+     * [onDispatchFutures] is the list of all [CompletableFuture]s that will complete because the [dispatchAll]
+     * method was invoked
+     *
+     * @return weather or not all futures gathered before [dispatchAll] were handled
+     */
+    fun onDispatchFuturesHandled(): Boolean =
+        onDispatchLoadCounter.get() == 0
+
+    /**
+     * Will signal if more dataLoaders where invoked during the [dispatchAll] invocation
+     * @return weather or not futures where loaded during [dispatchAll]
+     */
+    fun dataLoadersInvokedAfterDispatch(): Boolean =
+        onLoadCounter.get() > 0
+
+    fun onDataLoaderPromiseDispatched() {
+        onLoadCounter.incrementAndGet()
+    }
+
+    fun onDataLoaderPromiseCompleted(result: Any?, t: Throwable?) {
+        onDispatchLoadCounter.decrementAndGet()
+    }
+}

executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/execution/AbstractSyncExecutionExhaustedInstrumentation.kt renamed to executions/graphql-kotlin-dataloader-instrumentation/src/main/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/GraphQLSyncExecutionExhaustedDataLoaderDispatcher.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 Expedia, Inc
+ * Copyright 2025 Expedia, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -41,7 +41,7 @@ import java.util.concurrent.ConcurrentHashMap
  * }
  * ```
  *
- * When the [ExecutionBatchState] will be considered exhausted will have this state:
+ * The [ExecutionInputState] will be considered exhausted will have this state:
  *
  * ```
  * {
@@ -61,7 +61,7 @@ import java.util.concurrent.ConcurrentHashMap
  *
  * once astronaut [DataFetcher] completes his value starting a new ExecutionStrategy for `astronaut` field for
  * the resolution of `id`, and `name`, once these fields are dispatched and completed
- * and exhaustion will be calculated again and the state of [ExecutionBatchState] will be the following:
+ * and exhaustion will be calculated again and the state of [ExecutionInputState] will be the following:
  *
  * ```
  * {
@@ -94,8 +94,9 @@ import java.util.concurrent.ConcurrentHashMap
  * }
  * ```
  */
-class ExecutionBatchState {
-
+class ExecutionInputState(
+    val executionInput: ExecutionInput,
+) {
     private val executionStrategiesState: ConcurrentHashMap<String, ExecutionStrategyState> = ConcurrentHashMap()
 
     /**
@@ -177,7 +178,7 @@ class ExecutionBatchState {
     }
 
     /**
-     * Recursively calculate the sync state of this [ExecutionBatchState],
+     * Recursively calculate the sync state of this [ExecutionInputState],
      * by traversing all [executionStrategiesState].
      *
      * @param executionStrategyPath the string representation of the executionStrategy which state will be calculated,