Merge pull request #207 from aPureBase/kdataloader-fix

Fix the dataloader dispatcher issue

Author: jeggy

build.gradle.kts

@@ -6,7 +6,7 @@ val sonatypeUsername: String? = System.getenv("sonatypeUsername")
 val sonatypePassword: String? = System.getenv("sonatypePassword")
 
 plugins {
-    id("com.github.ben-manes.versions") version "0.39.0"
+    id("com.github.ben-manes.versions") version "0.44.0"
     id("io.codearte.nexus-staging") version "0.30.0"
     id("de.marcphilipp.nexus-publish") version "0.4.0"
     jacoco

gradle.properties

@@ -3,18 +3,18 @@ version=0.18.1
 
 # Dependencies
 coroutine_version=1.6.4
-jackson_version=2.13.4
-caffeine_version=3.1.1
-serialization_version=1.4.0
+jackson_version=2.14.1
+caffeine_version=3.1.2
+serialization_version=1.4.1
 kDataLoader_version=0.5.1
 deferredJsonBuilder_version=1.0.0
-ktor_version=2.1.2
+ktor_version=2.2.2
 
 # Test-Dependencies
 kotlin_html_version=0.7.5
 netty_version=4.1.82.Final
-junit_version=5.9.0
-kluent_version=1.68
+junit_version=5.9.2
+kluent_version=1.72
 hamcrest_version=2.2
 
 
@@ -25,7 +25,7 @@ coverallsGradlePlugin_version=2.8.4
 jacoco_version=0.8.5
 
 # Example Dependencies
-logback_version=1.2.1
+logback_version=1.4.5
 exposed_version=0.32.1
 h2_version=1.4.200
 hikari_version=4.0.3

kgraphql-example/build.gradle.kts

@@ -1,8 +1,8 @@
 plugins {
     base
     application
-    kotlin("jvm") version "1.7.10"
-    id("org.jetbrains.dokka") version "1.7.10"
+    kotlin("jvm") version "1.8.0"
+    id("org.jetbrains.dokka") version "1.7.20"
     signing
 }
 

kgraphql-ktor/build.gradle.kts

@@ -1,8 +1,8 @@
 plugins {
     base
-    kotlin("jvm") version "1.6.20"
-    kotlin("plugin.serialization") version "1.6.20"
-    id("org.jetbrains.dokka") version "1.4.32"
+    kotlin("jvm") version "1.8.0"
+    kotlin("plugin.serialization") version "1.8.0"
+    id("org.jetbrains.dokka") version "1.7.20"
     signing
 }
 

kgraphql/build.gradle.kts

@@ -1,8 +1,8 @@
 
 plugins {
     base
-    kotlin("jvm") version "1.7.10"
-    id("org.jetbrains.dokka") version "1.7.10"
+    kotlin("jvm") version "1.8.0"
+    id("org.jetbrains.dokka") version "1.7.20"
     signing
 }
 
@@ -26,14 +26,16 @@ dependencies {
     implementation(kotlin("reflect"))
 
     implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:$coroutine_version")
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-jdk8:$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")
+
+//    api("de.nidomiro:KDataLoader:$kDataLoader_version")
 
 
     testImplementation("io.netty:netty-all:$netty_version")

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

@@ -87,7 +87,6 @@ class DataLoaderPropertyDSL<T, K, R>(
                 { TimedAutoDispatcherDataLoaderOptions() },
                 mapOf(),
                 dataLoader!!,
-                null,
             )
         )
     }

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

@@ -16,9 +16,7 @@ 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.CompletableDeferred
-import kotlinx.coroutines.Deferred
-import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.*
 import kotlinx.serialization.json.*
 import nidomiro.kdataloader.DataLoader
 import kotlin.reflect.KProperty1
@@ -34,25 +32,24 @@ class DataLoaderPreparedRequestExecutor(val schema: DefaultSchema) : RequestExec
         val loaders: Map<Field.DataLoader<*, *, *>, DataLoader<Any?, *>>
     )
 
-    private suspend fun ExecutionPlan.constructLoaders(): Map<Field.DataLoader<*, *, *>, DataLoader<Any?, *>> {
+    private suspend fun ExecutionPlan.constructLoaders(): Map<Field.DataLoader<*, *, *>, DataLoader<Any?, *>> = coroutineScope {
         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?, *>
+                            loaders[ex.field] = ex.field.loader.constructNew(coroutineContext.job) as DataLoader<Any?, *>
                         }
                     }
                 }
             }
         }
         operations.look()
-        return loaders
+        loaders
     }
 
     private suspend fun <T> DeferredJsonMap.writeOperation(

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

@@ -292,7 +292,7 @@ class ParallelRequestExecutor(val schema: DefaultSchema) : RequestExecutor {
         )
 
         // as this isn't the DataLoaderPreparedRequestExecutor. We'll use this instant workaround instead.
-        val loader = field.loader.constructNew() as DataLoader<Any?, Any?>
+        val loader = field.loader.constructNew(null) as DataLoader<Any?, Any?>
         val value = loader.loadAsync(preparedValue)
         loader.dispatch()
 

kgraphql/src/main/kotlin/nidomiro/kdataloader/BatchMode.kt

@@ -0,0 +1,13 @@
+package nidomiro.kdataloader
+
+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/nidomiro/kdataloader/Cache.kt

@@ -0,0 +1,30 @@
+package nidomiro.kdataloader
+
+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/nidomiro/kdataloader/CoroutineMapCache.kt

@@ -0,0 +1,54 @@
+package nidomiro.kdataloader
+
+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()
+        }
+
+
+}

kgraphql/src/main/kotlin/nidomiro/kdataloader/DataLoader.kt

@@ -0,0 +1,83 @@
+package nidomiro.kdataloader
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.joinAll
+import kotlin.jvm.JvmName
+import nidomiro.kdataloader.statistics.DataLoaderStatistics
+
+typealias BatchLoader<K, R> = suspend (ids: List<K>) -> List<ExecutionResult<R>>
+
+interface DataLoader<K, R> {
+
+    val options: DataLoaderOptions<K, R>
+
+    /**
+     * Loads the value for the given Key.
+     * The returned [Deferred] completes with the finish of [dispatch] in case of [DataLoaderOptions.batchLoadEnabled] = true.
+     * If [DataLoaderOptions.batchLoadEnabled] = false it calls the BatchLoader immediately and returns the retrieved value.
+     */
+    suspend fun loadAsync(key: K): Deferred<R>
+
+    /**
+     * The same as [loadAsync] but for multiple Keys at once.
+     */
+    suspend fun loadManyAsync(vararg keys: K): Deferred<List<R>>
+
+    /**
+     * Executes all stored requests via the given batchLoader.
+     * After this function finishes all [Deferred] created before are completed.
+     */
+    suspend fun dispatch()
+
+    /**
+     * Removes the value of the given Key from the cache
+     */
+    suspend fun clear(key: K)
+
+    /**
+     * Removes all values from the cache
+     */
+    suspend fun clearAll()
+
+    /**
+     * Primes the cache with the given values.
+     * After priming the [BatchLoader] will not be called with this key.
+     */
+    suspend fun prime(key: K, value: R)
+
+    /**
+     * Primes the cache with the given [Throwable].
+     * After priming the [BatchLoader] will not be called with this key, if [DataLoaderOptions.cacheExceptions] = true.
+     */
+    suspend fun prime(key: K, value: Throwable)
+
+
+    /**
+     * Returns a snapshot of the statistics at the point of calling
+     */
+    suspend fun createStatisticsSnapshot(): DataLoaderStatistics
+
+}
+
+/**
+ * @see DataLoader.prime(K, R)
+ */
+suspend fun <K, R> SimpleDataLoaderImpl<K, R>.prime(cacheEntry: Pair<K, R>) {
+    prime(cacheEntry.first, cacheEntry.second)
+}
+
+/**
+ * @see DataLoader.prime(K, Throwable)
+ */
+@JvmName("primeFailure")
+suspend fun <K, R> SimpleDataLoaderImpl<K, R>.prime(cacheEntry: Pair<K, Throwable>) {
+    prime(cacheEntry.first, cacheEntry.second)
+}
+
+internal suspend fun <K, R> DataLoader<K, R>.prime(key: K, value: ExecutionResult<R>) =
+    when (value) {
+        is ExecutionResult.Success -> prime(key, value.value)
+        is ExecutionResult.Failure -> prime(key, value.throwable)
+    }

kgraphql/src/main/kotlin/nidomiro/kdataloader/DataLoaderOptions.kt

@@ -0,0 +1,18 @@
+package nidomiro.kdataloader
+
+open class DataLoaderOptions<K, R>(
+    /**
+     * The cache implementation
+     */
+    val cache: Cache<K, R>? = CoroutineMapCache(),
+
+    /**
+     * Cache Exceptional States?
+     */
+    val cacheExceptions: Boolean = true,
+
+    /**
+     * The batch-mode
+     */
+    val batchMode: BatchMode = BatchMode.LoadInBatch()
+)

kgraphql/src/main/kotlin/nidomiro/kdataloader/DefaultLoaderQueueImpl.kt

@@ -0,0 +1,27 @@
+package nidomiro.kdataloader
+
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+
+class DefaultLoaderQueueImpl<K, V> : LoaderQueue<K, V> {
+
+    private val mutex = Mutex()
+    private var queue: MutableList<LoaderQueueEntry<K, CompletableDeferred<V>>> = mutableListOf()
+
+
+    override suspend fun enqueue(key: K, deferred: CompletableDeferred<V>) {
+        mutex.withLock {
+            queue.add(LoaderQueueEntry(key, deferred))
+        }
+    }
+
+    override suspend fun getAllItemsAsList(): List<LoaderQueueEntry<K, CompletableDeferred<V>>> =
+        mutex.withLock {
+            val currentQueue = queue
+            queue = mutableListOf()
+            return@withLock currentQueue
+        }
+
+
+}

kgraphql/src/main/kotlin/nidomiro/kdataloader/ExecutionResult.kt

@@ -0,0 +1,6 @@
+package nidomiro.kdataloader
+
+sealed class ExecutionResult<out T> {
+    data class Success<out T>(val value: T) : ExecutionResult<T>()
+    data class Failure(val throwable: Throwable) : ExecutionResult<Nothing>()
+}