Update DataLoaderExecutor to be auto timed

Author: jeggy

gradle.properties

@@ -6,7 +6,8 @@ coroutine_version=1.5.0
 jackson_version=2.12.3
 caffeine_version=3.0.2
 serialization_version=1.2.1
-kDataLoader_version=0.3.0
+kDataLoader_version=0.4.0
+deferredJsonBuilder_version=1.0.0
 ktor_version=1.6.0
 
 # Test-Dependencies

kgraphql/build.gradle.kts

@@ -8,6 +8,7 @@ plugins {
 
 val caffeine_version: String by project
 val kDataLoader_version: String by project
+val deferredJsonBuilder_version: String by project
 val kotlin_version: String by project
 val serialization_version: String by project
 val coroutine_version: String by project
@@ -25,11 +26,14 @@ dependencies {
     implementation(kotlin("reflect"))
 
     implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:$coroutine_version")
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:$serialization_version") // JVM dependency
+
     implementation("com.fasterxml.jackson.core:jackson-databind:$jackson_version")
     implementation("com.fasterxml.jackson.module:jackson-module-kotlin:$jackson_version")
+
     implementation("com.github.ben-manes.caffeine:caffeine:$caffeine_version")
+    implementation("com.apurebase:DeferredJsonBuilder:$deferredJsonBuilder_version")
     api("de.nidomiro:KDataLoader:$kDataLoader_version")
-    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:$serialization_version") // JVM dependency
 
 
     testImplementation("io.netty:netty-all:$netty_version")

kgraphql/src/main/kotlin/com/apurebase/kgraphql/schema/dsl/DataLoaderPropertyDSL.kt

@@ -5,7 +5,8 @@ import com.apurebase.kgraphql.schema.model.FunctionWrapper
 import com.apurebase.kgraphql.schema.model.InputValueDef
 import com.apurebase.kgraphql.schema.model.PropertyDef
 import nidomiro.kdataloader.BatchLoader
-import nidomiro.kdataloader.dsl.dataLoaderFactory
+import nidomiro.kdataloader.TimedAutoDispatcherDataLoaderOptions
+import nidomiro.kdataloader.factories.TimedAutoDispatcherDataLoaderFactory
 import kotlin.reflect.KType
 
 class DataLoaderPropertyDSL<T, K, R>(
@@ -80,7 +81,12 @@ class DataLoaderPropertyDSL<T, K, R>(
             inputValues = inputValues,
             returnType = returnType,
             prepare = prepareWrapper!!,
-            loader = dataLoaderFactory(dataLoader!!)
+            loader = TimedAutoDispatcherDataLoaderFactory(
+                { TimedAutoDispatcherDataLoaderOptions() },
+                mapOf(),
+                dataLoader!!,
+                null,
+            )
         )
     }
 

kgraphql/src/main/kotlin/com/apurebase/kgraphql/schema/execution/DataLoaderPreparedRequestExecutor.kt

@@ -1,5 +1,7 @@
 package com.apurebase.kgraphql.schema.execution
 
+import com.apurebase.deferredJson.DeferredJsonMap
+import com.apurebase.deferredJson.deferredJsonBuilder
 import com.apurebase.kgraphql.Context
 import com.apurebase.kgraphql.ExecutionException
 import com.apurebase.kgraphql.GraphQLError
@@ -14,14 +16,11 @@ import com.apurebase.kgraphql.schema.scalar.serializeScalar
 import com.apurebase.kgraphql.schema.structure.Field
 import com.apurebase.kgraphql.schema.structure.InputValue
 import com.apurebase.kgraphql.schema.structure.Type
-import kotlinx.coroutines.*
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.coroutineScope
 import kotlinx.serialization.json.*
 import nidomiro.kdataloader.DataLoader
-import nidomiro.kdataloader.factories.DataLoaderFactory
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicLong
 import kotlin.reflect.KProperty1
 
 
@@ -32,46 +31,30 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
     inner class ExecutionContext(
         val variables: Variables,
         val requestContext: Context,
-        private val dataCounters: ConcurrentHashMap<DataLoader<*, *>, AtomicLong> = ConcurrentHashMap(),
-        private val loaders: ConcurrentHashMap<DataLoaderFactory<Any?, *>, DataLoader<Any?, *>> = ConcurrentHashMap()
-    ) : Mutex by Mutex() {
-
-
-        suspend fun get(loader: DataLoader<*, *>): Long = withLock {
-            dataCounters[loader]?.get() ?: throw IllegalArgumentException("Something went wrong with execution")
-        }
-        suspend fun add(loader: DataLoader<*, *>, check: Pair<*, String>, count: Long) = withLock {
-            if (dataCounters[loader] == null) {
-                /** TODO: There shouldn't be a need for Atomic here as we are using [withLock] */
-                dataCounters[loader] = AtomicLong(count)
-            } else {
-                val counter = dataCounters[loader]!!
-                counter.getAndUpdate {
-                    it + count
+        val loaders: Map<Field.DataLoader<*, *, *>, DataLoader<Any?, *>>
+    )
+
+    private suspend fun ExecutionPlan.constructLoaders(): Map<Field.DataLoader<*, *, *>, DataLoader<Any?, *>> {
+        val loaders = mutableMapOf<Field.DataLoader<*, *, *>, DataLoader<Any?, *>>()
+
+        suspend fun Collection<Execution>.look() {
+            forEach { ex ->
+                ex.selectionNode
+                when (ex) {
+                    is Execution.Fragment -> ex.elements.look()
+                    is Execution.Node -> {
+                        ex.children.look()
+                        if (ex.field is Field.DataLoader<*, *, *>) {
+                            loaders[ex.field] = ex.field.loader.constructNew() as DataLoader<Any?, *>
+                        }
+                    }
                 }
             }
         }
-
-        suspend fun load(builder: DeferredJsonMap, loaderfactory: DataLoaderFactory<Any?, Any?>, node: Execution.Node, preparedValue: Any?): Deferred<Any?> {
-            val loader = loaders.getOrPut(loaderfactory) { loaderfactory.constructNew() }
-            add(loader, preparedValue to node.selectionNode.fullPath, 1) // parentCount)
-
-            val value = loader.loadAsync(preparedValue)
-
-            builder.deferredLaunch {
-                val count = get(loader)
-                val stats = loader.createStatisticsSnapshot()
-                if (stats.objectsRequested >= count) {
-                    loader.dispatch()
-                } // else if (stats.objectsRequested > count) throw TODO("This should never happen!!!")
-            }
-
-
-            return value
-        }
+        operations.look()
+        return loaders
     }
 
-
     private suspend fun <T> DeferredJsonMap.writeOperation(
         ctx: ExecutionContext,
         node: Execution.Node,
@@ -212,7 +195,11 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
                     }
                 }
             }
-        } else {
+        } else if (expectedType.kind == TypeKind.UNION) return handleFragment(
+            ctx,
+            value,
+            container.elements.first { expectedType.name == expectedType.name } as Execution.Fragment
+        ) else {
             throw IllegalStateException("fragments can be specified on object types, interfaces, and unions")
         }
     }
@@ -326,7 +313,7 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
         ) // ?: TODO("Nullable prepare functions isn't supported")
 
 
-        val value = ctx.load(this, field.loader as DataLoaderFactory<Any?, Any?>, node, preparedValue)
+        val value = ctx.loaders[field]!!.loadAsync(preparedValue)
 
 
         applyKeyToElement(ctx, value, node, field.returnType, parentCount)
@@ -360,10 +347,11 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
     }
 
     override suspend fun suspendExecute(plan: ExecutionPlan, variables: VariablesJson, context: Context) = coroutineScope {
-        deferredJsonBuilder(timeout = plan.options.timeout ?: schema.configuration.timeout) {
+        val result = deferredJsonBuilder(timeout = plan.options.timeout ?: schema.configuration.timeout) {
             val ctx = ExecutionContext(
                 Variables(schema, variables, plan.firstOrNull { it.variables != null }?.variables),
-                context
+                context,
+                plan.constructLoaders(),
             )
 
 
@@ -372,7 +360,10 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
                     if (shouldInclude(ctx, node)) writeOperation(ctx, node, node.field as Field.Function<*, *>)
                 }
             }
-        }.toString()
+            ctx.loaders.values.map { it.dispatch() }
+        }
+
+        result.await().toString()
     }
 
     private fun createNullNode(node: Execution.Node, returnType: Type): JsonNull = if (returnType !is Type.NonNull) {