return ChainedInstrumentation when registry present in context

Author: agairol

src/main/java/graphql/servlet/OsgiGraphQLServlet.java

@@ -1,5 +1,6 @@
 package graphql.servlet;
 
+import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
 import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
@@ -207,7 +208,13 @@ protected ExecutionStrategyProvider getExecutionStrategyProvider() {
     @Override
     protected Instrumentation getInstrumentation(GraphQLContext context) {
         return context.getDataLoaderRegistry()
-                .map(DataLoaderDispatcherInstrumentation::new)
+                .map(registry -> {
+                    List<Instrumentation> instrumentations = new ArrayList<>();
+                    instrumentations.add(this.instrumentationProvider.getInstrumentation());
+                    instrumentations.add(new DataLoaderDispatcherInstrumentation(registry));
+
+                    return new ChainedInstrumentation(instrumentations);
+                })
                 .map(Instrumentation.class::cast)
                 .orElse(instrumentationProvider.getInstrumentation());
     }

src/main/java/graphql/servlet/SimpleGraphQLServlet.java

@@ -1,24 +1,28 @@
 package graphql.servlet;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
 import graphql.execution.ExecutionStrategy;
+import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
 import graphql.execution.instrumentation.SimpleInstrumentation;
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation;
 import graphql.execution.preparsed.NoOpPreparsedDocumentProvider;
 import graphql.execution.preparsed.PreparsedDocumentProvider;
 import graphql.schema.GraphQLSchema;
 
+
+
 /**
  * @author Andrew Potter
  */
 public class SimpleGraphQLServlet extends GraphQLServlet {
-	
+
 
     /**
      * @deprecated use {@link #builder(GraphQLSchema)} instead.
@@ -52,7 +56,7 @@ public SimpleGraphQLServlet(final GraphQLSchema schema, ExecutionStrategyProvide
         this(new DefaultGraphQLSchemaProvider(schema), executionStrategyProvider, objectMapperConfigurer, listeners, instrumentation, errorHandler, contextBuilder, rootObjectBuilder, preparsedDocumentProvider,false);
     }
 
-    
+
     /**
      * @deprecated use {@link #builder(GraphQLSchemaProvider)} instead.
      */
@@ -189,7 +193,7 @@ public Builder withListeners(List<GraphQLServletListener> listeners) {
             this.listeners = listeners;
             return this;
         }
-        
+
         public Builder withAsyncServletMode(boolean value) {
         	this.asyncServletMode=value;
         	return this;
@@ -223,7 +227,13 @@ protected ExecutionStrategyProvider getExecutionStrategyProvider() {
     @Override
     protected Instrumentation getInstrumentation(GraphQLContext context) {
         return context.getDataLoaderRegistry()
-                .map(DataLoaderDispatcherInstrumentation::new)
+                .map(registry -> {
+                    List<Instrumentation> instrumentations = new ArrayList<>();
+                    instrumentations.add(this.instrumentation);
+                    instrumentations.add(new DataLoaderDispatcherInstrumentation(registry));
+
+                    return new ChainedInstrumentation(instrumentations);
+                })
                 .map(Instrumentation.class::cast)
                 .orElse(this.instrumentation);
     }

src/test/groovy/graphql/servlet/GraphQLServletSpec.groovy

@@ -3,6 +3,7 @@ package graphql.servlet
 import com.fasterxml.jackson.databind.ObjectMapper
 import graphql.Scalars
 import graphql.execution.ExecutionTypeInfo
+import graphql.execution.instrumentation.ChainedInstrumentation
 import graphql.execution.instrumentation.Instrumentation
 import graphql.execution.instrumentation.dataloader.DataLoaderDispatcherInstrumentation
 import graphql.schema.DataFetcher
@@ -882,7 +883,7 @@ class GraphQLServletSpec extends Specification {
 
     }
 
-    def "getInstrumentation returns the DataLoaderDispatcherInstrumentation if DataLoader provided in context"() {
+    def "getInstrumentation returns the ChainedInstrumentation if DataLoader provided in context"() {
 
         setup:
             Instrumentation servletInstrumentation = Mock()
@@ -896,8 +897,7 @@ class GraphQLServletSpec extends Specification {
         when:
             Instrumentation actualInstrumentation = simpleGraphQLServlet.getInstrumentation(context)
         then:
-            actualInstrumentation instanceof DataLoaderDispatcherInstrumentation
-            actualInstrumentation != servletInstrumentation;
-
+            actualInstrumentation instanceof ChainedInstrumentation
+            actualInstrumentation != servletInstrumentation
     }
 }