Repository: netflix/dgs-framework
Issue: bug: `graphql-dgs-spring-boot-micrometer` Dataloader instrumentation does not work with `DgsDataLoaderRegistryConsumer` loaders

## Expected behavior

Instrumented Dataloaders implementing `DgsDataLoaderRegistryConsumer` in order to get a `DataLoaderRegistry` should have their `setDataLoaderRegistry` method called prior to execution.

## Actual behavior

`setDataLoaderRegistry` is never invoked on instrumented dataloaders. 

## Steps to reproduce

1. Create a dataloader that implements `DgsDataLoaderRegistryConsumer` and `setDataLoaderRegistry`.
```java
public class SomeDataLoader
    implements MappedBatchLoaderWithContext<String, ...>, DgsDataLoaderRegistryConsumer {
  // ...
  @Override
  public void setDataLoaderRegistry(DataLoaderRegistry dataLoaderRegistry) {
    this.dataLoaderRegistry = dataLoaderRegistry;
  }
}
```
2. Add in the instrumentation package into the pom.xml (I tested on 4.9.17, but after reading a bit of the code it appears to be an issue with future versions as well)
```xml
        <dependency>
            <groupId>com.netflix.graphql.dgs</groupId>
            <artifactId>graphql-dgs-spring-boot-micrometer</artifactId>
        </dependency>
```
3. Drop a breakpoint into the dataloader's load method (or otherwise add some application logic that requires the registry) and observe that `setDataLoaderRegistry` was never called.

## Mitigation

If you turn off DataLoader instrumentation with `dgs-graphql.data-loader-instrumentation.enabled=false` the data loader registry is set on the registry consumer as expected -- but you do lose out on dataloader metrics entirely

Changed files:
graphql-dgs-spring-boot-micrometer/src/test/kotlin/com/netflix/graphql/dgs/metrics/micrometer/MicrometerServletSmokeTest.kt
graphql-dgs-spring-boot-micrometer/src/main/kotlin/com/netflix/graphql/dgs/metrics/micrometer/dataloader/DgsDataLoaderInstrumentationProvider.kt
graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt

diff --git a/graphql-dgs-spring-boot-micrometer/src/main/kotlin/com/netflix/graphql/dgs/metrics/micrometer/dataloader/DgsDataLoaderInstrumentationProvider.kt b/graphql-dgs-spring-boot-micrometer/src/main/kotlin/com/netflix/graphql/dgs/metrics/micrometer/dataloader/DgsDataLoaderInstrumentationProvider.kt
index 3c92537f..688bd0f1 100644
--- a/graphql-dgs-spring-boot-micrometer/src/main/kotlin/com/netflix/graphql/dgs/metrics/micrometer/dataloader/DgsDataLoaderInstrumentationProvider.kt
+++ b/graphql-dgs-spring-boot-micrometer/src/main/kotlin/com/netflix/graphql/dgs/metrics/micrometer/dataloader/DgsDataLoaderInstrumentationProvider.kt
@@ -29,7 +29,8 @@ class DgsDataLoaderInstrumentationProvider(
                     .to(BatchLoaderInterceptor(original, name, meterRegistrySupplier.get()))
             ByteBuddy()
                 .subclass(BatchLoader::class.java)
-                .method(ElementMatchers.named("load")).intercept(withBinders)
+                .method(ElementMatchers.named("load"))
+                .intercept(withBinders)
                 .make()
                 .load(javaClass.classLoader)
                 .loaded
@@ -44,7 +45,8 @@ class DgsDataLoaderInstrumentationProvider(
 
             ByteBuddy()
                 .subclass(BatchLoaderWithContext::class.java)
-                .method(ElementMatchers.named("load")).intercept(withBinders)
+                .method(ElementMatchers.named("load"))
+                .intercept(withBinders)
                 .make()
                 .load(javaClass.classLoader)
                 .loaded
@@ -60,7 +62,8 @@ class DgsDataLoaderInstrumentationProvider(
 
             ByteBuddy()
                 .subclass(MappedBatchLoader::class.java)
-                .method(ElementMatchers.named("load")).intercept(withBinders)
+                .method(ElementMatchers.named("load"))
+                .intercept(withBinders)
                 .make()
                 .load(javaClass.classLoader)
                 .loaded
@@ -77,7 +80,8 @@ class DgsDataLoaderInstrumentationProvider(
                 .to(MappedBatchLoaderWithContextInterceptor(original, name, meterRegistrySupplier.get()))
             ByteBuddy()
                 .subclass(MappedBatchLoaderWithContext::class.java)
-                .method(ElementMatchers.named("load")).intercept(withBinders)
+                .method(ElementMatchers.named("load"))
+                .intercept(withBinders)
                 .make()
                 .load(javaClass.classLoader)
                 .loaded
diff --git a/graphql-dgs-spring-boot-micrometer/src/test/kotlin/com/netflix/graphql/dgs/metrics/micrometer/MicrometerServletSmokeTest.kt b/graphql-dgs-spring-boot-micrometer/src/test/kotlin/com/netflix/graphql/dgs/metrics/micrometer/MicrometerServletSmokeTest.kt
index 80e60c47..0321e1d2 100644
--- a/graphql-dgs-spring-boot-micrometer/src/test/kotlin/com/netflix/graphql/dgs/metrics/micrometer/MicrometerServletSmokeTest.kt
+++ b/graphql-dgs-spring-boot-micrometer/src/test/kotlin/com/netflix/graphql/dgs/metrics/micrometer/MicrometerServletSmokeTest.kt
@@ -19,6 +19,7 @@ package com.netflix.graphql.dgs.metrics.micrometer
 import com.netflix.graphql.dgs.DgsComponent
 import com.netflix.graphql.dgs.DgsData
 import com.netflix.graphql.dgs.DgsDataLoader
+import com.netflix.graphql.dgs.DgsDataLoaderRegistryConsumer
 import com.netflix.graphql.dgs.DgsEnableDataFetcherInstrumentation
 import com.netflix.graphql.dgs.DgsQuery
 import com.netflix.graphql.dgs.DgsTypeDefinitionRegistry
@@ -45,6 +46,7 @@ import io.micrometer.core.instrument.Tags
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry
 import org.assertj.core.api.Assertions.assertThat
 import org.dataloader.BatchLoader
+import org.dataloader.DataLoaderRegistry
 import org.dataloader.MappedBatchLoader
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -890,13 +892,19 @@ class MicrometerServletSmokeTest {
             @DgsDataLoader(name = "reverser")
             class ReverseStringDataLoader(
                 @Qualifier("dataLoaderTaskExecutor") private val dataLoaderTaskExecutor: Executor
-            ) : MappedBatchLoader<String, String> {
+            ) : MappedBatchLoader<String, String>, DgsDataLoaderRegistryConsumer {
+
+                var dataLoaderRegistry: Optional<DataLoaderRegistry> = Optional.empty()
+
                 override fun load(keys: Set<String>): CompletionStage<Map<String, String>> {
                     return CompletableFuture.supplyAsync(
                         { keys.associateWith { it.reversed() } },
                         dataLoaderTaskExecutor
                     )
                 }
+                override fun setDataLoaderRegistry(dataLoaderRegistry: DataLoaderRegistry) {
+                    this.dataLoaderRegistry = Optional.of(dataLoaderRegistry)
+                }
             }
         }
 
diff --git a/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt b/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt
index 2df4d209..7b0ebbb4 100644
--- a/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt
+++ b/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsDataLoaderProvider.kt
@@ -146,11 +146,11 @@ class DgsDataLoaderProvider(private val applicationContext: ApplicationContext)
     ): DataLoader<*, *> {
         val options = dataLoaderOptions(dgsDataLoader)
 
-        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
-        if (extendedBatchLoader is DgsDataLoaderRegistryConsumer) {
-            extendedBatchLoader.setDataLoaderRegistry(dataLoaderRegistry)
+        if (batchLoader is DgsDataLoaderRegistryConsumer) {
+            batchLoader.setDataLoaderRegistry(dataLoaderRegistry)
         }
 
+        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
         return DataLoaderFactory.newDataLoader(extendedBatchLoader, options)
     }
 
@@ -161,10 +161,10 @@ class DgsDataLoaderProvider(private val applicationContext: ApplicationContext)
     ): DataLoader<*, *> {
         val options = dataLoaderOptions(dgsDataLoader)
 
-        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
-        if (extendedBatchLoader is DgsDataLoaderRegistryConsumer) {
-            extendedBatchLoader.setDataLoaderRegistry(dataLoaderRegistry)
+        if (batchLoader is DgsDataLoaderRegistryConsumer) {
+            batchLoader.setDataLoaderRegistry(dataLoaderRegistry)
         }
+        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
 
         return DataLoaderFactory.newMappedDataLoader(extendedBatchLoader, options)
     }
@@ -178,11 +178,11 @@ class DgsDataLoaderProvider(private val applicationContext: ApplicationContext)
         val options = dataLoaderOptions(dgsDataLoader)
             .setBatchLoaderContextProvider(supplier::get)
 
-        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
-        if (extendedBatchLoader is DgsDataLoaderRegistryConsumer) {
-            extendedBatchLoader.setDataLoaderRegistry(dataLoaderRegistry)
+        if (batchLoader is DgsDataLoaderRegistryConsumer) {
+            batchLoader.setDataLoaderRegistry(dataLoaderRegistry)
         }
 
+        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
         return DataLoaderFactory.newDataLoader(extendedBatchLoader, options)
     }
 
@@ -195,11 +195,11 @@ class DgsDataLoaderProvider(private val applicationContext: ApplicationContext)
         val options = dataLoaderOptions(dgsDataLoader)
             .setBatchLoaderContextProvider(supplier::get)
 
-        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
-        if (extendedBatchLoader is DgsDataLoaderRegistryConsumer) {
-            extendedBatchLoader.setDataLoaderRegistry(dataLoaderRegistry)
+        if (batchLoader is DgsDataLoaderRegistryConsumer) {
+            batchLoader.setDataLoaderRegistry(dataLoaderRegistry)
         }
 
+        val extendedBatchLoader = wrappedDataLoader(batchLoader, dgsDataLoader.name)
         return DataLoaderFactory.newMappedDataLoader(extendedBatchLoader, options)
     }