Merge pull request #149 from aPureBase/temp-solution-for-data-loader-issues

Use new timed auto dispatcher for dataloaders

Author: jeggy

gradle.properties

@@ -1,5 +1,5 @@
 # KGraphQL version
-version=0.17.8
+version=0.17.9-alpha
 
 # Dependencies
 coroutine_version=1.3.9

kgraphql-ktor/build.gradle.kts

@@ -1,7 +1,7 @@
 plugins {
     base
-    kotlin("jvm") version "1.4.31"
-    kotlin("plugin.serialization") version "1.4.31"
+    kotlin("jvm") version "1.5.10"
+    kotlin("plugin.serialization") version "1.5.0"
     id("org.jetbrains.dokka") version "0.10.1"
     signing
 }

kgraphql/build.gradle.kts

@@ -1,7 +1,7 @@
 
 plugins {
     base
-    kotlin("jvm") version "1.4.31"
+    kotlin("jvm") version "1.5.10"
     id("org.jetbrains.dokka") version "0.10.1"
     signing
 }
@@ -28,7 +28,7 @@ dependencies {
     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")
-    api("de.nidomiro:KDataLoader:$kDataLoader_version")
+//    api("de.nidomiro:KDataLoader:$kDataLoader_version")
     implementation("org.jetbrains.kotlinx:kotlinx-serialization-core:$serialization_version") // JVM dependency
 
 

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

@@ -4,8 +4,10 @@ import com.apurebase.kgraphql.Context
 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 com.apurebase.kgraphql.test.BatchLoader
+import com.apurebase.kgraphql.test.TimedAutoDispatcherDataLoaderOptions
+import com.apurebase.kgraphql.test.dsl.dataLoaderFactory
+import com.apurebase.kgraphql.test.factories.TimedAutoDispatcherDataLoaderFactory
 import kotlin.reflect.KType
 
 class DataLoaderPropertyDSL<T, K, R>(
@@ -80,7 +82,11 @@ class DataLoaderPropertyDSL<T, K, R>(
             inputValues = inputValues,
             returnType = returnType,
             prepare = prepareWrapper!!,
-            loader = dataLoaderFactory(dataLoader!!)
+            loader = TimedAutoDispatcherDataLoaderFactory(
+                { TimedAutoDispatcherDataLoaderOptions() },
+                mapOf(),
+                dataLoader!!
+            )
         )
     }
 

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

@@ -15,13 +15,8 @@ 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.serialization.json.*
-import nidomiro.kdataloader.DataLoader
-import nidomiro.kdataloader.factories.DataLoaderFactory
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicLong
+import com.apurebase.kgraphql.test.DataLoader
 import kotlin.reflect.KProperty1
 
 
@@ -32,43 +27,29 @@ 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(job: Job): 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(job) 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
     }
 
 
@@ -325,8 +306,7 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
             ctx = ctx
         ) // ?: 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,19 +340,23 @@ 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(job),
             )
 
-
             "data" toDeferredObj {
                 plan.forEach { node ->
                     if (shouldInclude(ctx, node)) writeOperation(ctx, node, node.field as Field.Function<*, *>)
                 }
             }
-        }.toString()
+            ctx.loaders.values.map { it.dispatch() }
+        }
+
+
+        result.toString()
     }
 
     private fun createNullNode(node: Execution.Node, returnType: Type): JsonNull = if (returnType !is Type.NonNull) {

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

@@ -18,7 +18,7 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory
 import com.fasterxml.jackson.databind.node.NullNode
 import com.fasterxml.jackson.databind.node.ObjectNode
 import kotlinx.coroutines.*
-import nidomiro.kdataloader.DataLoader
+import com.apurebase.kgraphql.test.DataLoader
 import kotlin.reflect.KProperty1
 
 

kgraphql/src/main/kotlin/com/apurebase/kgraphql/schema/model/PropertyDef.kt

@@ -1,7 +1,7 @@
 package com.apurebase.kgraphql.schema.model
 
 import com.apurebase.kgraphql.Context
-import nidomiro.kdataloader.BatchLoader
+import com.apurebase.kgraphql.test.factories.DataLoaderFactory
 import kotlin.reflect.KProperty1
 import kotlin.reflect.KType
 
@@ -29,7 +29,7 @@ interface PropertyDef<T> : Depreciable, DescribedDef {
      */
     open class DataLoadedFunction<T, K, R>(
         override val name: String,
-        val loader: nidomiro.kdataloader.factories.DataLoaderFactory<K, R>,
+        val loader: DataLoaderFactory<K, R>,
         val prepare: FunctionWrapper<K>,
         val returnType: KType,
         override val description: String? = null,

kgraphql/src/main/kotlin/com/apurebase/kgraphql/schema/structure/Field.kt

@@ -3,7 +3,7 @@ package com.apurebase.kgraphql.schema.structure
 import com.apurebase.kgraphql.Context
 import com.apurebase.kgraphql.schema.introspection.*
 import com.apurebase.kgraphql.schema.model.*
-import nidomiro.kdataloader.DataLoader
+import com.apurebase.kgraphql.test.factories.DataLoaderFactory
 import kotlin.reflect.full.findAnnotation
 
 
@@ -44,7 +44,7 @@ sealed class Field : __Field {
 
     class DataLoader<T, K, R>(
         val kql: PropertyDef.DataLoadedFunction<T, K, R>,
-        val loader: nidomiro.kdataloader.factories.DataLoaderFactory<K, R>,
+        val loader: DataLoaderFactory<K, R>,
         override val returnType: Type,
         override val arguments: List<InputValue<*>>
     ): Field() {

kgraphql/src/main/kotlin/com/apurebase/kgraphql/test/BatchMode.kt

@@ -0,0 +1,13 @@
+package com.apurebase.kgraphql.test
+
+sealed class BatchMode {
+    /**
+     * Load data in batches of [batchSize]
+     */
+    data class LoadInBatch(val batchSize: Int? = null) : BatchMode()
+
+    /**
+     * Load everything immediately
+     */
+    object LoadImmediately : BatchMode()
+}

kgraphql/src/main/kotlin/com/apurebase/kgraphql/test/Cache.kt

@@ -0,0 +1,30 @@
+package com.apurebase.kgraphql.test
+
+import kotlinx.coroutines.CompletableDeferred
+
+/**
+ * A "Threadsafe" Cache
+ * (Coroutine-Save)
+ */
+interface Cache<K, V> {
+
+    suspend fun store(key: K, value: CompletableDeferred<V>): CompletableDeferred<V>
+
+    suspend fun get(key: K): CompletableDeferred<V>?
+
+    suspend fun getOrCreate(
+        key: K,
+        generator: suspend (key: K) -> CompletableDeferred<V>,
+        callOnCacheHit: suspend () -> Unit
+    ): CompletableDeferred<V>
+
+    suspend fun clear(key: K): CompletableDeferred<V>?
+
+    suspend fun clear()
+}
+
+suspend fun <K, V> Cache<K, V>.getOrCreate(
+    key: K,
+    generator: suspend (key: K) -> CompletableDeferred<V>
+): CompletableDeferred<V> =
+    getOrCreate(key, generator, {})

kgraphql/src/main/kotlin/com/apurebase/kgraphql/test/CoroutineMapCache.kt

@@ -0,0 +1,54 @@
+package com.apurebase.kgraphql.test
+
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+
+class CoroutineMapCache<K, V>(
+    private val cacheMap: MutableMap<K, CompletableDeferred<V>> = mutableMapOf()
+) : Cache<K, V> {
+    private val mutex = Mutex()
+
+
+    override suspend fun store(key: K, value: CompletableDeferred<V>): CompletableDeferred<V> {
+        mutex.withLock {
+            cacheMap[key] = value
+        }
+        return value
+    }
+
+    override suspend fun get(key: K): CompletableDeferred<V>? =
+        mutex.withLock {
+            cacheMap[key]
+        }
+
+    override suspend fun getOrCreate(
+        key: K,
+        generator: suspend (key: K) -> CompletableDeferred<V>,
+        callOnCacheHit: suspend () -> Unit
+    ): CompletableDeferred<V> =
+        mutex.withLock {
+            val currentVal = cacheMap[key]
+            if (currentVal == null) {
+                val generated = generator(key)
+                cacheMap[key] = generated
+                return@withLock generated
+            } else {
+                callOnCacheHit()
+                return@withLock currentVal
+            }
+        }
+
+
+    override suspend fun clear(key: K): CompletableDeferred<V>? =
+        mutex.withLock {
+            cacheMap.remove(key)
+        }
+
+    override suspend fun clear() =
+        mutex.withLock {
+            cacheMap.clear()
+        }
+
+
+}