feat: initial commit

Author: Samuel Vazquez

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

@@ -0,0 +1,38 @@
+/*
+ * 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
+
+import com.expediagroup.graphql.dataloader.KotlinDataLoaderInstrumentation
+import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
+import org.dataloader.instrumentation.DataLoaderInstrumentationContext
+
+internal class DataLoaderInstrumentationContext(
+    private val syncExecutionExhaustedState: SyncExecutionExhaustedState
+): DataLoaderInstrumentationContext<Any?> {
+    override fun onDispatched() {
+        syncExecutionExhaustedState.onDataLoaderPromiseDispatched()
+    }
+    override fun onCompleted(result: Any?, t: Throwable?) {
+        syncExecutionExhaustedState.onDataLoaderPromiseCompleted()
+    }
+}
+
+class DataLoaderInstrumentationForSyncExecutionExhausted(
+    syncExecutionExhaustedState: SyncExecutionExhaustedState
+) : KotlinDataLoaderInstrumentation(
+    DataLoaderInstrumentationContext(syncExecutionExhaustedState)
+)

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

@@ -39,10 +39,7 @@ import java.util.concurrent.atomic.AtomicInteger
  * Orchestrate the [ExecutionBatchState] of all [ExecutionInput] sharing the same [GraphQLContext],
  * when a certain state is reached will invoke [OnSyncExecutionExhaustedCallback]
  */
-class SyncExecutionExhaustedState(
-    totalOperations: Int,
-    private val dataLoaderRegistry: KotlinDataLoaderRegistry
-) {
+class SyncExecutionExhaustedState(totalOperations: Int) {
     private val totalExecutions: AtomicInteger = AtomicInteger(totalOperations)
     val executions = ConcurrentHashMap<ExecutionId, ExecutionBatchState>()
 
@@ -128,7 +125,7 @@ class SyncExecutionExhaustedState(
 
         return object : FieldFetchingInstrumentationContext {
             override fun onFetchedValue(fetchedValue: Any?) {
-                println("${fieldName}:$path")
+                //println("${fieldName}:$path")
                 executions.computeIfPresent(executionId) { _, executionState ->
                     executionState.fieldToDispatchedState(field, fieldExecutionStrategyPath, fieldGraphQLType, fetchedValue)
                     executionState
@@ -141,7 +138,7 @@ class SyncExecutionExhaustedState(
 
             override fun onCompleted(result: Any?, t: Throwable?) {
                 executions.computeIfPresent(executionId) { _, executionState ->
-                    println("${fieldName}:$path")
+                    println("completed: ${fieldName}:$path")
                     executionState.fieldToCompletedState(field, fieldExecutionStrategyPath, result)
                     executionState
                 }
@@ -156,6 +153,14 @@ class SyncExecutionExhaustedState(
         }
     }
 
+    fun onDataLoaderPromiseDispatched() {
+        dataLoaderRegistry.onDataLoaderPromiseDispatched()
+    }
+
+    fun onDataLoaderPromiseCompleted() {
+        dataLoaderRegistry.onDataLoaderPromiseCompleted()
+    }
+
     /**
      * execute a given [predicate] when all [ExecutionInput] sharing a [GraphQLContext] exhausted their execution.
      * A Synchronous Execution is considered Exhausted when all [DataFetcher]s of all paths were executed up until

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

@@ -32,7 +32,6 @@ import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.P
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.PlanetServiceRequest
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Mission
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Nasa
-import com.expediagroup.graphql.dataloader.DataLoaderDependantsStateInstrumentation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import graphql.ExecutionInput
 import graphql.ExecutionResult
@@ -43,13 +42,11 @@ import graphql.schema.idl.RuntimeWiring
 import graphql.schema.idl.SchemaGenerator
 import graphql.schema.idl.SchemaParser
 import graphql.schema.idl.TypeRuntimeWiring
-import io.mockk.mockk
 import io.mockk.spyk
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.runBlocking
-import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation
 
 object AstronautGraphQL {
     private val schema = """

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

@@ -24,7 +24,6 @@ import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.P
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Product
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductDetails
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductSummary
-import com.expediagroup.graphql.dataloader.DataLoaderDependantsStateInstrumentation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import graphql.ExecutionInput
 import graphql.ExecutionResult
@@ -36,13 +35,11 @@ import graphql.schema.idl.RuntimeWiring
 import graphql.schema.idl.SchemaGenerator
 import graphql.schema.idl.SchemaParser
 import graphql.schema.idl.TypeRuntimeWiring
-import io.mockk.mockk
 import io.mockk.spyk
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.runBlocking
-import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation
 import java.util.concurrent.CompletableFuture
 
 object ProductGraphQL {

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentationTest.kt renamed to executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/syncexhaustion/DataLoaderInstrumentationForSyncExecutionExhaustedTest.kt

@@ -27,7 +27,7 @@ import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
 
-class DataLoaderSyncExecutionExhaustedInstrumentationTest {
+class DataLoaderInstrumentationForSyncExecutionExhaustedTest {
     private val dataLoaderSyncExecutionExhaustedInstrumentation = DataLoaderSyncExecutionExhaustedInstrumentation()
     private val astronautGraphQL = AstronautGraphQL.builder
         .instrumentation(dataLoaderSyncExecutionExhaustedInstrumentation)
@@ -71,7 +71,7 @@ class DataLoaderSyncExecutionExhaustedInstrumentationTest {
         assertEquals(1, missionStatistics?.batchInvokeCount)
         assertEquals(2, missionStatistics?.batchLoadCount)
 
-        verify(exactly = 2) {
+        verify(exactly = 1) {
             kotlinDataLoaderRegistry.dispatchAll()
         }
     }

executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/DataLoaderDependantsStateInstrumentation.kt renamed to executions/graphql-kotlin-dataloader/src/main/kotlin/com/expediagroup/graphql/dataloader/KotlinDataLoaderInstrumentation.kt

@@ -16,29 +16,16 @@
 
 package com.expediagroup.graphql.dataloader
 
-import graphql.GraphQLContext
 import org.dataloader.DataLoader
 import org.dataloader.instrumentation.DataLoaderInstrumentation
 import org.dataloader.instrumentation.DataLoaderInstrumentationContext
 
-class DataLoaderDependantsStateInstrumentation(
-    private val graphQLContext: GraphQLContext
+open class KotlinDataLoaderInstrumentation(
+    private val context: DataLoaderInstrumentationContext<Any?>
 ) : DataLoaderInstrumentation {
     override fun beginLoad(
         dataLoader: DataLoader<*, *>,
         key: Any,
         loadContext: Any?
-    ): DataLoaderInstrumentationContext<Any> =
-        object : DataLoaderInstrumentationContext<Any> {
-            override fun onDispatched() {
-                graphQLContext
-                    .get<KotlinDataLoaderRegistry>(KotlinDataLoaderRegistry::class)
-                    ?.onDataLoaderPromiseDispatched()
-            }
-            override fun onCompleted(result: Any?, t: Throwable?) {
-                graphQLContext
-                    .get<KotlinDataLoaderRegistry>(KotlinDataLoaderRegistry::class)
-                    ?.onDataLoaderPromiseCompleted()
-            }
-        }
+    ): DataLoaderInstrumentationContext<Any?> = context
 }

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

@@ -25,31 +25,25 @@ import org.dataloader.instrumentation.DataLoaderInstrumentation
  * Generates a [KotlinDataLoaderRegistry] with the configuration provided by all [KotlinDataLoader]s.
  */
 class KotlinDataLoaderRegistryFactory(
-    private val dataLoaders: List<KotlinDataLoader<*, *>>,
-    private val dataLoaderInstrumentations: List<DataLoaderInstrumentation>
+    private val dataLoaders: List<KotlinDataLoader<*, *>>
 ) {
-    constructor(): this(emptyList(), emptyList())
-    constructor(dataLoaders: List<KotlinDataLoader<*, *>>): this(dataLoaders, emptyList())
     /**
      * Generate [KotlinDataLoaderRegistry] to be used for GraphQL request execution.
      */
-    fun generate(graphQLContext: GraphQLContext): KotlinDataLoaderRegistry {
+    fun generate(
+        graphQLContext: GraphQLContext,
+        dataLoaderInstrumentation: DataLoaderInstrumentation? = null,
+    ): KotlinDataLoaderRegistry {
         val builder = DataLoaderRegistry.newRegistry()
-        builder.instrumentation(
-            ChainedDataLoaderInstrumentation(
-                dataLoaderInstrumentations.toMutableList().also {
-                    it.add(DataLoaderDependantsStateInstrumentation(graphQLContext))
-                }
-            )
-        )
+        dataLoaderInstrumentation?.let {
+            builder.instrumentation(dataLoaderInstrumentation)
+        }
         dataLoaders.forEach { kotlinDataLoader ->
             builder.register(
                 kotlinDataLoader.dataLoaderName,
                 kotlinDataLoader.getDataLoader(graphQLContext)
             )
         }
-        return KotlinDataLoaderRegistry(builder.build()).also {
-            graphQLContext.put(KotlinDataLoaderRegistry::class, it)
-        }
+        return KotlinDataLoaderRegistry(builder.build())
     }
 }

servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandler.kt

@@ -18,6 +18,7 @@ package com.expediagroup.graphql.server.execution
 
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistry
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
+import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.DataLoaderInstrumentationForSyncExecutionExhausted
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.DataLoaderSyncExecutionExhaustedInstrumentation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import com.expediagroup.graphql.generator.extensions.plus
@@ -46,6 +47,7 @@ import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.supervisorScope
+import org.dataloader.instrumentation.DataLoaderInstrumentation
 
 open class GraphQLRequestHandler(
     private val graphQL: GraphQL,
@@ -73,21 +75,25 @@ open class GraphQLRequestHandler(
      */
     open suspend fun executeRequest(
         graphQLRequest: GraphQLServerRequest,
-        graphQLContext: GraphQLContext = GraphQLContext.of(emptyMap<Any, Any>())
+        graphQLContext: GraphQLContext = GraphQLContext.getDefault()
     ): GraphQLServerResponse {
-        val dataLoaderRegistry = dataLoaderRegistryFactory.generate(graphQLContext)
         return when (graphQLRequest) {
             is GraphQLRequest -> {
-                val batchGraphQLContext = graphQLContext + getBatchContext(1, dataLoaderRegistry)
+                val (batchContext, dataLoaderInstrumentation) = getBatchContext(1)
+                val batchGraphQLContext = graphQLContext + batchContext
+                val dataLoaderRegistry = dataLoaderRegistryFactory.generate(batchGraphQLContext, dataLoaderInstrumentation)
                 execute(graphQLRequest, batchGraphQLContext, dataLoaderRegistry)
             }
-
             is GraphQLBatchRequest -> {
                 if (graphQLRequest.containsMutation()) {
-                    val batchGraphQLContext = graphQLContext + getBatchContext(1, dataLoaderRegistry)
+                    val (batchContext, dataLoaderInstrumentation) = getBatchContext(1)
+                    val batchGraphQLContext = graphQLContext + batchContext
+                    val dataLoaderRegistry = dataLoaderRegistryFactory.generate(batchGraphQLContext, dataLoaderInstrumentation)
                     executeSequentially(graphQLRequest, batchGraphQLContext, dataLoaderRegistry)
                 } else {
-                    val batchGraphQLContext = graphQLContext + getBatchContext(graphQLRequest.requests.size, dataLoaderRegistry)
+                    val (batchContext, dataLoaderInstrumentation) = getBatchContext(graphQLRequest.requests.size)
+                    val batchGraphQLContext = graphQLContext + batchContext
+                    val dataLoaderRegistry = dataLoaderRegistryFactory.generate(batchGraphQLContext, dataLoaderInstrumentation)
                     executeConcurrently(graphQLRequest, batchGraphQLContext, dataLoaderRegistry)
                 }
             }
@@ -134,17 +140,18 @@ open class GraphQLRequestHandler(
         return GraphQLBatchResponse(responses)
     }
 
-    private fun getBatchContext(
-        batchSize: Int,
-        dataLoaderRegistry: KotlinDataLoaderRegistry
-    ): Map<*, Any> =
-        when (batchDataLoaderInstrumentationType) {
-            DataLoaderSyncExecutionExhaustedInstrumentation::class.java -> mapOf(
-                KotlinDataLoaderRegistry::class to dataLoaderRegistry,
-                SyncExecutionExhaustedState::class to SyncExecutionExhaustedState(batchSize, dataLoaderRegistry)
-            )
-            else -> emptyMap<Any, Any>()
+    private fun getBatchContext(batchSize: Int): Pair<GraphQLContext, DataLoaderInstrumentation?> {
+        return when (batchDataLoaderInstrumentationType) {
+            DataLoaderSyncExecutionExhaustedInstrumentation::class.java -> {
+                val syncExecutionExhaustedState = SyncExecutionExhaustedState(batchSize)
+                GraphQLContext.of(
+                    mapOf(SyncExecutionExhaustedState::class to syncExecutionExhaustedState)
+                ) to DataLoaderInstrumentationForSyncExecutionExhausted(syncExecutionExhaustedState)
+            }
+
+            else -> GraphQLContext.getDefault() to null
         }
+    }
 
     /**
      * Execute a GraphQL subscription operation in a non-blocking fashion.

servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/GraphQLExecutionConfiguration.kt

@@ -19,7 +19,7 @@ package com.expediagroup.graphql.server.spring
 import com.expediagroup.graphql.generator.execution.KotlinDataFetcherFactoryProvider
 import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
 import com.expediagroup.graphql.dataloader.KotlinDataLoader
-import com.expediagroup.graphql.dataloader.DataLoaderDependantsStateInstrumentation
+import com.expediagroup.graphql.dataloader.KotlinDataLoaderInstrumentation
 import com.expediagroup.graphql.server.spring.execution.SpringKotlinDataFetcherFactoryProvider
 import graphql.execution.DataFetcherExceptionHandler
 import graphql.execution.SimpleDataFetcherExceptionHandler
@@ -63,7 +63,7 @@ class GraphQLExecutionConfiguration {
             instrumentations.addAll(it)
         }
         if (config.batching.enabled) {
-            instrumentations.add(DataLoaderDependantsStateInstrumentation())
+            instrumentations.add(KotlinDataLoaderInstrumentation())
         }
         return KotlinDataLoaderRegistryFactory(
             dataLoaders.orElse(emptyList()),