Merge pull request #1818 from Netflix/ejd--new-dataloader-hooks

add new hooks for instrumenting dataloaders

graphql-dgs-spring-boot-oss-autoconfigure/src/main/kotlin/com/netflix/graphql/dgs/autoconfig/DgsAutoConfiguration.kt

@@ -17,6 +17,8 @@
 package com.netflix.graphql.dgs.autoconfig
 
 import com.netflix.graphql.dgs.DataLoaderInstrumentationExtensionProvider
+import com.netflix.graphql.dgs.DgsDataLoaderCustomizer
+import com.netflix.graphql.dgs.DgsDataLoaderInstrumentation
 import com.netflix.graphql.dgs.DgsDataLoaderOptionsProvider
 import com.netflix.graphql.dgs.DgsDefaultPreparsedDocumentProvider
 import com.netflix.graphql.dgs.DgsFederationResolver
@@ -166,20 +168,42 @@ open class DgsAutoConfiguration(
         return Executors.newSingleThreadScheduledExecutor()
     }
 
+    @Bean
+    @ConditionalOnProperty(
+        prefix = "$AUTO_CONF_PREFIX.convertAllDataLoadersToWithContext",
+        name = ["enabled"],
+        havingValue = "true",
+        matchIfMissing = true
+    )
+    @Order(0)
+    open fun dgsWrapWithContextDataLoaderCustomizer(): DgsWrapWithContextDataLoaderCustomizer {
+        return DgsWrapWithContextDataLoaderCustomizer()
+    }
+
+    @Bean
+    @Order(100)
+    open fun dgsDataLoaderInstrumentationDataLoaderCustomizer(
+        instrumentations: List<DgsDataLoaderInstrumentation>
+    ): DgsDataLoaderInstrumentationDataLoaderCustomizer {
+        return DgsDataLoaderInstrumentationDataLoaderCustomizer(instrumentations)
+    }
+
     @Bean
     open fun dgsDataLoaderProvider(
         applicationContext: ApplicationContext,
         dataloaderOptionProvider: DgsDataLoaderOptionsProvider,
         @Qualifier("dgsScheduledExecutorService") dgsScheduledExecutorService: ScheduledExecutorService,
-        extensionProviders: List<DataLoaderInstrumentationExtensionProvider>
+        extensionProviders: List<DataLoaderInstrumentationExtensionProvider>,
+        customizers: List<DgsDataLoaderCustomizer>
     ): DgsDataLoaderProvider {
         return DgsDataLoaderProvider(
             applicationContext = applicationContext,
             extensionProviders = extensionProviders,
             dataLoaderOptionsProvider = dataloaderOptionProvider,
             scheduledExecutorService = dgsScheduledExecutorService,
             scheduleDuration = dataloaderConfigProps.scheduleDuration,
-            enableTickerMode = dataloaderConfigProps.tickerModeEnabled
+            enableTickerMode = dataloaderConfigProps.tickerModeEnabled,
+            customizers = customizers
         )
     }
 

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/DgsDataLoaderCustomizer.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2024 Netflix, 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
+ *
+ *    http://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.netflix.graphql.dgs
+
+import org.dataloader.BatchLoader
+import org.dataloader.BatchLoaderWithContext
+import org.dataloader.MappedBatchLoader
+import org.dataloader.MappedBatchLoaderWithContext
+
+/**
+ * Beans that implement this interface will be called during the component scan
+ * that finds defined @DgsDataLoaders and allows the modification or wrapping of
+ * each DataLoader.
+ *
+ * While this hook appears very similar to [com.netflix.graphql.dgs.DataLoaderInstrumentationExtensionProvider]
+ * there are two important differences:
+ * - DgsDataLoaderCustomizers are called when scanning for @DgsDataLoaders, whereas the
+ *   DataLoaderInstrumentationExtensionProviders are called once per request
+ * - DgsDataLoaderCustomizers are afforded the opportunity to change the type of each DataLoader
+ *   as it is being initially registered. Most notably, this allows for converting BatchLoader and MappedBatchLoader
+ *   into their "WithContext" versions, which can simplify certain types of DataLoader instrumentation.
+ */
+interface DgsDataLoaderCustomizer {
+    fun provide(original: BatchLoader<*, *>, name: String): Any {
+        return original
+    }
+    fun provide(original: BatchLoaderWithContext<*, *>, name: String): Any {
+        return original
+    }
+    fun provide(original: MappedBatchLoader<*, *>, name: String): Any {
+        return original
+    }
+    fun provide(original: MappedBatchLoaderWithContext<*, *>, name: String): Any {
+        return original
+    }
+}

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/DgsDataLoaderInstrumentation.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2024 Netflix, 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
+ *
+ *    http://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.netflix.graphql.dgs
+
+import org.dataloader.BatchLoaderEnvironment
+
+/**
+ * Provides a simple interface to define instrumentation around the dispatch of data loaders.
+ *
+ * Requires [com.netflix.graphql.dgs.internal.DgsWrapWithContextDataLoaderCustomizer] to be enabled,
+ * as DgsDataLoaderInstrumentation only provides hooks for the WithContext versions of batch loaders.
+ */
+interface DgsDataLoaderInstrumentation {
+    /**
+     * onDispatch will run before the data loader itself is actually called.
+     *
+     * @param name the name of the data loader
+     * @param keys the list of keys dispatched to the data loader
+     * @param batchLoaderEnvironment the batchLoaderEnvironment for the current execution
+     *
+     * @return context object that also contains the other hooks
+     */
+    fun onDispatch(name: String, keys: List<Any>, batchLoaderEnvironment: BatchLoaderEnvironment): DgsDataLoaderInstrumentationContext
+}

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/DgsDataLoaderInstrumentationContext.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2024 Netflix, 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
+ *
+ *    http://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.netflix.graphql.dgs
+
+interface DgsDataLoaderInstrumentationContext {
+    /**
+     * onComplete will run in a whenComplete attached to the data loader's returned CompletableFuture.
+     *
+     * @param result the actual results of the data loader. Will be a Map or List depending on the type of data loader.
+     * @param exception any exception thrown by the data loader
+     *
+     * This means it will run on one of two possible threads:
+     *  - the thread associated with the data loader's CompletableFuture
+     *  - the thread that actually dispatched the data loader
+     */
+    fun onComplete(result: Any?, exception: Any?)
+}

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/exceptions/DgsDataLoaderInstrumentationException.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2024 Netflix, 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
+ *
+ *    http://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.netflix.graphql.dgs.exceptions
+
+class DgsDataLoaderInstrumentationException(name: String) : RuntimeException("data loader `$name` is not a MappedBatchLoaderWithContext or BatchLoaderWithContext. Is dgs.graphql.convertAllDataLoadersToWithContext.enabled set to false?")

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderInstrumentationDataLoaderCustomizer.kt

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2024 Netflix, 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
+ *
+ *    http://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.netflix.graphql.dgs.internal
+
+import com.netflix.graphql.dgs.DgsDataLoaderCustomizer
+import com.netflix.graphql.dgs.DgsDataLoaderInstrumentation
+import com.netflix.graphql.dgs.exceptions.DgsDataLoaderInstrumentationException
+import org.dataloader.BatchLoader
+import org.dataloader.BatchLoaderEnvironment
+import org.dataloader.BatchLoaderWithContext
+import org.dataloader.MappedBatchLoader
+import org.dataloader.MappedBatchLoaderWithContext
+import java.util.concurrent.CompletionStage
+
+class DgsDataLoaderInstrumentationDataLoaderCustomizer(
+    private val instrumentations: List<DgsDataLoaderInstrumentation>
+) : DgsDataLoaderCustomizer {
+    override fun provide(original: BatchLoader<*, *>, name: String): Any {
+        throw DgsDataLoaderInstrumentationException(name)
+    }
+
+    override fun provide(original: BatchLoaderWithContext<*, *>, name: String): Any {
+        return BatchLoaderWithContextInstrumentationDriver(original, name, instrumentations)
+    }
+
+    override fun provide(original: MappedBatchLoader<*, *>, name: String): Any {
+        throw DgsDataLoaderInstrumentationException(name)
+    }
+
+    override fun provide(original: MappedBatchLoaderWithContext<*, *>, name: String): Any {
+        return MappedBatchLoaderWithContextInstrumentationDriver(original, name, instrumentations)
+    }
+}
+
+internal class BatchLoaderWithContextInstrumentationDriver<K : Any, V>(
+    private val original: BatchLoaderWithContext<K, V>,
+    private val name: String,
+    private val instrumentations: List<DgsDataLoaderInstrumentation>
+) : BatchLoaderWithContext<K, V> {
+    override fun load(keys: List<K>, environment: BatchLoaderEnvironment): CompletionStage<List<V>> {
+        val contexts = instrumentations.map { it.onDispatch(name, keys, environment) }
+        val future = original.load(keys, environment)
+
+        return future.whenComplete { result, exception ->
+            try {
+                contexts.asReversed().forEach { c ->
+                    c.onComplete(result, exception)
+                }
+            } catch (_: Throwable) {
+            }
+        }
+    }
+}
+
+internal class MappedBatchLoaderWithContextInstrumentationDriver<K : Any, V>(
+    private val original: MappedBatchLoaderWithContext<K, V>,
+    private val name: String,
+    private val instrumentations: List<DgsDataLoaderInstrumentation>
+) : MappedBatchLoaderWithContext<K, V> {
+    override fun load(keys: Set<K>, environment: BatchLoaderEnvironment): CompletionStage<Map<K, V>> {
+        val keysList = keys.toList()
+        val contexts = instrumentations.map { it.onDispatch(name, keysList, environment) }
+
+        val future = original.load(keys, environment)
+
+        return future.whenComplete { result, exception ->
+            try {
+                contexts.asReversed().forEach { c ->
+                    c.onComplete(result, exception)
+                }
+            } catch (_: Throwable) {
+            }
+        }
+    }
+}

graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt

@@ -50,6 +50,7 @@ import kotlin.system.measureTimeMillis
 class DgsDataLoaderProvider(
     private val applicationContext: ApplicationContext,
     private val extensionProviders: List<DataLoaderInstrumentationExtensionProvider> = listOf(),
+    private val customizers: List<DgsDataLoaderCustomizer> = listOf(),
     private val dataLoaderOptionsProvider: DgsDataLoaderOptionsProvider = DefaultDataLoaderOptionsProvider(),
     private val scheduledExecutorService: ScheduledExecutorService = Executors.newSingleThreadScheduledExecutor(),
     private val scheduleDuration: Duration = Duration.ofMillis(10),
@@ -107,11 +108,11 @@ class DgsDataLoaderProvider(
                     }
 
                     fun <T : Any> createHolder(t: T): LoaderHolder<T> = LoaderHolder(t, annotation, annotation.name)
-                    when (val get = field.get(dgsComponent)) {
-                        is BatchLoader<*, *> -> batchLoaders.add(createHolder(get))
-                        is BatchLoaderWithContext<*, *> -> batchLoadersWithContext.add(createHolder(get))
-                        is MappedBatchLoader<*, *> -> mappedBatchLoaders.add(createHolder(get))
-                        is MappedBatchLoaderWithContext<*, *> -> mappedBatchLoadersWithContext.add(createHolder(get))
+                    when (val dataLoader = runCustomizers(field.get(dgsComponent), annotation.name, dgsComponent::class.java)) {
+                        is BatchLoader<*, *> -> batchLoaders.add(createHolder(dataLoader))
+                        is BatchLoaderWithContext<*, *> -> batchLoadersWithContext.add(createHolder(dataLoader))
+                        is MappedBatchLoader<*, *> -> mappedBatchLoaders.add(createHolder(dataLoader))
+                        is MappedBatchLoaderWithContext<*, *> -> mappedBatchLoadersWithContext.add(createHolder(dataLoader))
                         else -> throw InvalidDataLoaderTypeException(dgsComponent::class.java)
                     }
                 }
@@ -137,17 +138,35 @@ class DgsDataLoaderProvider(
     }
 
     private fun <T : Any> addDataLoaders(dgsComponent: T, targetClass: Class<*>, annotation: DgsDataLoader, dispatchPredicate: DispatchPredicate?) {
+        val name = DataLoaderNameUtil.getDataLoaderName(targetClass, annotation)
         fun <T : Any> createHolder(t: T): LoaderHolder<T> =
             LoaderHolder(t, annotation, DataLoaderNameUtil.getDataLoaderName(targetClass, annotation), dispatchPredicate)
-        when (dgsComponent) {
-            is BatchLoader<*, *> -> batchLoaders.add(createHolder(dgsComponent))
-            is BatchLoaderWithContext<*, *> -> batchLoadersWithContext.add(createHolder(dgsComponent))
-            is MappedBatchLoader<*, *> -> mappedBatchLoaders.add(createHolder(dgsComponent))
-            is MappedBatchLoaderWithContext<*, *> -> mappedBatchLoadersWithContext.add(createHolder(dgsComponent))
+
+        when (val dataLoader = runCustomizers(dgsComponent, name, dgsComponent::class.java)) {
+            is BatchLoader<*, *> -> batchLoaders.add(createHolder(dataLoader))
+            is BatchLoaderWithContext<*, *> -> batchLoadersWithContext.add(createHolder(dataLoader))
+            is MappedBatchLoader<*, *> -> mappedBatchLoaders.add(createHolder(dataLoader))
+            is MappedBatchLoaderWithContext<*, *> -> mappedBatchLoadersWithContext.add(createHolder(dataLoader))
             else -> throw InvalidDataLoaderTypeException(dgsComponent::class.java)
         }
     }
 
+    private fun runCustomizers(originalDataLoader: Any, name: String, dgsComponentClass: Class<*>): Any {
+        var dataLoader = originalDataLoader
+
+        customizers.forEach {
+            dataLoader = when (dataLoader) {
+                is BatchLoader<*, *> -> it.provide(dataLoader as BatchLoader<*, *>, name)
+                is BatchLoaderWithContext<*, *> -> it.provide(dataLoader as BatchLoaderWithContext<*, *>, name)
+                is MappedBatchLoader<*, *> -> it.provide(dataLoader as MappedBatchLoader<*, *>, name)
+                is MappedBatchLoaderWithContext<*, *> -> it.provide(dataLoader as MappedBatchLoaderWithContext<*, *>, name)
+                else -> throw InvalidDataLoaderTypeException(dgsComponentClass)
+            }
+        }
+
+        return dataLoader
+    }
+
     private fun createDataLoader(
         batchLoader: BatchLoader<*, *>,
         dgsDataLoader: DgsDataLoader,