feat: add comments

Author: Samuel Vazquez

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

@@ -39,7 +39,7 @@ fun <V> CompletableFuture<V>.dispatchIfNeeded(
         .get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
         ?: throw MissingInstrumentationStateException()
 
-    if (syncExecutionExhaustedState.dataLoadersInvokedAfterDispatch() && syncExecutionExhaustedState.allSyncExecutionsExhausted()) {
+    if (syncExecutionExhaustedState.dataLoadersLoadInvokedAfterDispatchAll() && syncExecutionExhaustedState.allSyncExecutionsExhausted()) {
         dataLoaderRegistry.dispatchAll()
     }
     return this

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

@@ -16,22 +16,26 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion
 
+import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.DataLoaderRegistryState
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import org.dataloader.DataLoader
 import org.dataloader.instrumentation.DataLoaderInstrumentation
 import org.dataloader.instrumentation.DataLoaderInstrumentationContext
 
+/**
+ * Custom [DataLoaderInstrumentation] implementation that helps to calculate the state of [DataLoader]s in the
+ * [DataLoaderRegistryState] that lives inside the [syncExecutionExhaustedState]
+ */
 class DataLoaderSyncExecutionExhaustedDataLoaderDispatcher(
     private val syncExecutionExhaustedState: SyncExecutionExhaustedState
-): DataLoaderInstrumentation {
-
+) : DataLoaderInstrumentation {
     private val contextForSyncExecutionExhausted: DataLoaderInstrumentationContext<Any?> =
-        object: DataLoaderInstrumentationContext<Any?> {
+        object : DataLoaderInstrumentationContext<Any?> {
             override fun onDispatched() {
-                syncExecutionExhaustedState.onDataLoaderPromiseDispatched()
+                syncExecutionExhaustedState.onDataLoaderLoadDispatched()
             }
             override fun onCompleted(result: Any?, t: Throwable?) {
-                syncExecutionExhaustedState.onDataLoaderPromiseCompleted(result, t)
+                syncExecutionExhaustedState.onDataLoaderLoadCompleted()
             }
         }
 

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

@@ -30,7 +30,6 @@ 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
 
 /**
@@ -74,12 +73,4 @@ class GraphQLSyncExecutionExhaustedDataLoaderDispatcher : SimplePerformantInstru
         parameters.executionContext.takeUnless(ExecutionContext::isMutation)
             ?.graphQLContext?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
             ?.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/state/DataLoaderDispatchState.kt

@@ -0,0 +1,72 @@
+/*
+ * 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 org.dataloader.DataLoader
+import org.dataloader.DataLoaderRegistry
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * Calculate the state of [DataLoader.load]s in a [DataLoaderRegistry]
+ */
+class DataLoaderRegistryState {
+    /**
+     * Count [DataLoader.load] invocations
+     */
+    private val loadCounter = AtomicInteger(0)
+
+    /**
+     * Snapshot of [loadCounter] when [DataLoaderRegistry.dispatchAll] is invoked,
+     * then on every load complete decrease it
+     */
+    private val onDispatchAllLoadCounter = AtomicInteger(0)
+
+    /**
+     * Take snapshot of [loadCounter] when [DataLoaderRegistry.dispatchAll] is invoked
+     */
+    fun takeSnapshot() {
+        onDispatchAllLoadCounter.set(loadCounter.get())
+        loadCounter.set(0)
+    }
+
+    /**
+     * @return if all [CompletableFuture]s returned by [DataLoader.load] were completed
+     */
+    fun onDispatchAllFuturesCompleted(): Boolean =
+        onDispatchAllLoadCounter.get() == 0
+
+    /**
+     * @return If more [DataLoader.load] where invoked after the [DataLoaderRegistry.dispatchAll] invocation
+     */
+    fun dataLoadersLoadInvokedAfterDispatchAll(): Boolean =
+        loadCounter.get() > 0
+
+    /**
+     * Increase [loadCounter] when [DataLoader.load] is invoked
+     */
+    fun onDataLoaderLoadDispatched() {
+        loadCounter.incrementAndGet()
+    }
+
+    /**
+     * Decrease [onDispatchAllLoadCounter] when [DataLoader.load] returned [CompletableFuture] completes
+     */
+    fun onDataLoaderLoadCompleted() {
+        onDispatchAllLoadCounter.decrementAndGet()
+    }
+}

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

@@ -29,6 +29,7 @@ import graphql.execution.instrumentation.parameters.InstrumentationExecutionPara
 import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
 import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
 import graphql.schema.DataFetcher
+import org.dataloader.DataLoader
 import org.dataloader.DataLoaderRegistry
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ConcurrentHashMap
@@ -42,9 +43,9 @@ class SyncExecutionExhaustedState(
     totalOperations: Int,
     private val dataLoaderRegistryProvider: () -> DataLoaderRegistry
 ) {
-    private val totalExecutions: AtomicInteger = AtomicInteger(totalOperations)
+    private val totalExecutions = AtomicInteger(totalOperations)
     private val executions = ConcurrentHashMap<ExecutionId, ExecutionInputState>()
-    private val dataLoadersDispatchState = DataLoaderDispatchState()
+    private val dataLoadersDispatchState = DataLoaderRegistryState()
 
     /**
      * Create the [ExecutionInputState] When a specific [ExecutionInput] starts his execution
@@ -107,7 +108,7 @@ class SyncExecutionExhaustedState(
     }
 
     /**
-     * This is called just before a field [DataFetcher] is invoked
+     * This is invoked just before a field [DataFetcher] is invoked
      *
      * @param parameters contains information of which field will start the fetching
      * @return a [InstrumentationContext] object that will be called back when the [DataFetcher]
@@ -153,16 +154,21 @@ class SyncExecutionExhaustedState(
         }
     }
 
-    fun dataLoadersInvokedAfterDispatch(): Boolean =
-        dataLoadersDispatchState.dataLoadersInvokedAfterDispatch()
+    fun dataLoadersLoadInvokedAfterDispatchAll(): Boolean =
+        dataLoadersDispatchState.dataLoadersLoadInvokedAfterDispatchAll()
 
-    fun onDataLoaderPromiseDispatched() {
-        dataLoadersDispatchState.onDataLoaderPromiseDispatched()
+    /**
+     * This is invoked right after a [DataLoader.load] was dispatched
+     */
+    fun onDataLoaderLoadDispatched() {
+        dataLoadersDispatchState.onDataLoaderLoadDispatched()
     }
 
-    fun onDataLoaderPromiseCompleted(result: Any?, t: Throwable?) {
-        dataLoadersDispatchState.onDataLoaderPromiseCompleted(result, t)
-        // println("completed dataloader promise: $result")
+    /**
+     * This is invoked right after a [DataLoader.load] was completed
+     */
+    fun onDataLoaderLoadCompleted() {
+        dataLoadersDispatchState.onDataLoaderLoadCompleted()
         if (allSyncExecutionsExhausted()) {
             dataLoaderRegistryProvider.invoke().dispatchAll()
         }
@@ -175,7 +181,7 @@ class SyncExecutionExhaustedState(
      */
     fun allSyncExecutionsExhausted(): Boolean =
         synchronized(executions) {
-            if (executions.size < totalExecutions.get() || !dataLoadersDispatchState.onDispatchFuturesHandled())
+            if (executions.size < totalExecutions.get() || !dataLoadersDispatchState.onDispatchAllFuturesCompleted())
                 return false
 
             if (executions.values.all(ExecutionInputState::isSyncExecutionExhausted)) {

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

@@ -140,7 +140,7 @@ object ProductGraphQL {
                         ExecutionInput
                             .newExecutionInput(query)
                             .dataLoaderRegistry(dataLoaderRegistry)
-                            .graphQLContext { it.of(graphQLContext)  }
+                            .graphQLContext { it.of(graphQLContext) }
                             .build()
                     ).await()
                 }

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt

@@ -26,7 +26,7 @@ import org.dataloader.instrumentation.DataLoaderInstrumentation
 class KotlinDataLoaderRegistryFactory(
     private val dataLoaders: List<KotlinDataLoader<*, *>>
 ) {
-    constructor(): this(emptyList())
+    constructor() : this(emptyList())
 
     /**
      * Generate [DataLoaderRegistry] to be used for GraphQL request execution.

executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderRegistryFactory.kt

@@ -20,7 +20,7 @@ import graphql.GraphQLContext
 import io.mockk.mockk
 import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
-import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation
+import org.dataloader.instrumentation.DataLoaderInstrumentation
 import org.junit.jupiter.api.Test
 import reactor.kotlin.core.publisher.toFlux
 import kotlin.test.assertEquals
@@ -29,7 +29,7 @@ import kotlin.test.assertTrue
 class KotlinDataLoaderRegistryFactoryTest {
     @Test
     fun `generate registry with empty list`() {
-        val registry = KotlinDataLoaderRegistryFactory().generate(mockk())
+        val registry = KotlinDataLoaderRegistryFactory().generate(mockk(relaxed = true))
         assertTrue(registry.dataLoaders.isEmpty())
     }
 
@@ -43,10 +43,16 @@ class KotlinDataLoaderRegistryFactoryTest {
                 }
         }
 
+        val customInstrumentation = object : DataLoaderInstrumentation {
+
+        }
         val registry = KotlinDataLoaderRegistryFactory(
             listOf(mockLoader)
-        ).generate(mockk())
+        ).generate(
+            mockk(relaxed = true),
+            customInstrumentation
+        )
         assertEquals(1, registry.dataLoaders.size)
-        assertTrue(registry.instrumentation is ChainedDataLoaderInstrumentation)
+        assertEquals(customInstrumentation, registry.instrumentation)
     }
 }