diff --git a/src/main/java/graphql/servlet/instrumentation/ConfigurableDispatchInstrumentation.java b/src/main/java/graphql/servlet/instrumentation/ConfigurableDispatchInstrumentation.java index 2817e258..d2c05164 100644 --- a/src/main/java/graphql/servlet/instrumentation/ConfigurableDispatchInstrumentation.java +++ b/src/main/java/graphql/servlet/instrumentation/ConfigurableDispatchInstrumentation.java @@ -162,19 +162,20 @@ public InstrumentationContext beginFieldFetch(InstrumentationFieldFetchP @Override public CompletableFuture instrumentExecutionResult(ExecutionResult executionResult, InstrumentationExecutionParameters parameters) { - if (!options.isIncludeStatistics()) { - return CompletableFuture.completedFuture(executionResult); - } DataLoaderDispatcherInstrumentationState state = parameters.getInstrumentationState(); - Map currentExt = executionResult.getExtensions(); - Map statsMap = new LinkedHashMap<>(currentExt == null ? Collections.emptyMap() : currentExt); - Map dataLoaderStats = buildStatsMap(state); - statsMap.put("dataloader", dataLoaderStats); state.getApproach().removeTracking(parameters.getExecutionInput().getExecutionId()); + if (!options.isIncludeStatistics()) { + return CompletableFuture.completedFuture(executionResult); + } else { + Map currentExt = executionResult.getExtensions(); + Map statsMap = new LinkedHashMap<>(currentExt == null ? Collections.emptyMap() : currentExt); + Map dataLoaderStats = buildStatsMap(state); + statsMap.put("dataloader", dataLoaderStats); - log.debug("Data loader stats : {}", dataLoaderStats); + log.debug("Data loader stats : {}", dataLoaderStats); - return CompletableFuture.completedFuture(new ExecutionResultImpl(executionResult.getData(), executionResult.getErrors(), statsMap)); + return CompletableFuture.completedFuture(new ExecutionResultImpl(executionResult.getData(), executionResult.getErrors(), statsMap)); + } } private Map buildStatsMap(DataLoaderDispatcherInstrumentationState state) { diff --git a/src/main/java/graphql/servlet/instrumentation/RequestStack.java b/src/main/java/graphql/servlet/instrumentation/RequestStack.java index 22c1c7a0..0b459dd1 100644 --- a/src/main/java/graphql/servlet/instrumentation/RequestStack.java +++ b/src/main/java/graphql/servlet/instrumentation/RequestStack.java @@ -117,7 +117,9 @@ public void setStatus(ExecutionId executionId, boolean toState) { throw new IllegalStateException( String.format("Can not set status for execution %s, it is not managed by this request stack", executionId)); } - status.put(executionId, toState); + if (toState) { + status.put(executionId, true); + } } /**