increase accuracy of engine running

andimarek · andimarek · commit 339f65788081 · 2025-04-02T21:30:45.000+10:00
diff --git a/src/main/java/graphql/execution/ExecutionContext.java b/src/main/java/graphql/execution/ExecutionContext.java
@@ -29,6 +29,7 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
@@ -370,31 +371,48 @@ public boolean isRunning() {
         return isRunning.get() > 0;
     }
 
-    public <T> T call(Supplier<T> callable) {
+    public void incrementRunning() {
         if (isRunning.incrementAndGet() == 1 && engineRunningObserver != null) {
             engineRunningObserver.runningStateChanged(executionId, graphQLContext, true);
         }
+    }
+
+    public void decrementRunning() {
+        if (isRunning.decrementAndGet() == 0 && engineRunningObserver != null) {
+            engineRunningObserver.runningStateChanged(executionId, graphQLContext, false);
+        }
+    }
+
+    public void incrementRunning(CompletableFuture<?> cf) {
+        cf.whenComplete((result, throwable) -> {
+            incrementRunning();
+        });
+    }
+
+    public void decrementRunning(CompletableFuture<?> cf) {
+        cf.whenComplete((result, throwable) -> {
+            decrementRunning();
+        });
+
+    }
+
+    public <T> T call(Supplier<T> callable) {
+        incrementRunning();
         try {
             return callable.get();
         } finally {
-            if (isRunning.decrementAndGet() == 0 && engineRunningObserver != null) {
-                engineRunningObserver.runningStateChanged(executionId, graphQLContext, false);
-            }
+            decrementRunning();
         }
     }
 
     public void run(Runnable runnable) {
-        if (isRunning.incrementAndGet() == 1 && engineRunningObserver != null) {
-            engineRunningObserver.runningStateChanged(executionId, graphQLContext, true);
-        }
+        incrementRunning();
         try {
             runnable.run();
         } finally {
-            if (isRunning.decrementAndGet() == 0 && engineRunningObserver != null) {
-                engineRunningObserver.runningStateChanged(executionId, graphQLContext, false);
-            }
-
+            decrementRunning();
         }
     }
 
+
 }
diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java
@@ -513,7 +513,8 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
         if (fetchedObject instanceof CompletableFuture) {
             @SuppressWarnings("unchecked")
             CompletableFuture<Object> fetchedValue = (CompletableFuture<Object>) fetchedObject;
-            return fetchedValue
+            executionContext.decrementRunning(fetchedValue);
+            CompletableFuture<FetchedValue> fetchedValueCF = fetchedValue
                     .handle((result, exception) -> executionContext.call(() -> {
                         fetchCtx.onCompleted(result, exception);
                         if (exception != null) {
@@ -526,6 +527,8 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
                     }))
                     .thenCompose(Function.identity())
                     .thenApply(result -> unboxPossibleDataFetcherResult(executionContext, parameters, result));
+            executionContext.incrementRunning(fetchedValue);
+            return fetchedValueCF;
         } else {
             fetchCtx.onCompleted(fetchedObject, null);
             return unboxPossibleDataFetcherResult(executionContext, parameters, fetchedObject);
diff --git a/src/test/groovy/graphql/EngineRunningTest.groovy b/src/test/groovy/graphql/EngineRunningTest.groovy
@@ -45,7 +45,7 @@ class EngineRunningTest extends Specification {
         states == [true, false]
     }
 
-    def "engine running state is observed with async datafetcher"() {
+    def "engine running state is observed with one async datafetcher"() {
         given:
         def sdl = '''
 
@@ -80,11 +80,62 @@ class EngineRunningTest extends Specification {
         cf.complete("world")
 
         then:
-        //TODO: why is that so much back and forth between true and false?
-        states == [true, false, true, false, true, false, true, false]
+        states == [true, false]
         er.get().data == [hello: "world"]
 
     }
 
+    def "engine running state is observed with two async datafetcher"() {
+        given:
+        def sdl = '''
+
+        type Query {
+          hello: String
+          hello2: String
+        }
+        '''
+        CompletableFuture cf1 = new CompletableFuture();
+        CompletableFuture cf2 = new CompletableFuture();
+        def df = { env ->
+            return cf1;
+        } as DataFetcher
+        def df2 = { env ->
+            return cf2
+        } as DataFetcher
+
+        def fetchers = ["Query": ["hello": df, "hello2": df2]]
+        def schema = TestUtil.schema(sdl, fetchers)
+        def graphQL = GraphQL.newGraphQL(schema).build()
+
+        def query = "{ hello hello2 }"
+        def ei = newExecutionInput(query).build()
+
+        List<Boolean> states = new CopyOnWriteArrayList<>();
+        ei.getGraphQLContext().put(ENGINE_RUNNING_OBSERVER_KEY, {
+            ExecutionId executionId, GraphQLContext context, boolean running ->
+                states.add(running)
+        } as EngineRunningObserver);
+
+        when:
+        def er = graphQL.executeAsync(ei)
+        then:
+        states == [true, false]
+
+        when:
+        states.clear();
+        cf1.complete("world")
+
+        then:
+        states == [true, false]
+
+        when:
+        states.clear();
+        cf2.complete("world2")
+
+        then:
+        states == [true, false]
+        er.get().data == [hello: "world", hello2: "world2"]
+    }
+
 
 }
