Skip to content

Commit 93b53e4

Browse files
bbakermanbbakerman
committed
Merge remote-tracking branch 'origin/extract-out-operation-directives' into extract-out-operation-directives
2 parents 601195a + 7b92be4 commit 93b53e4

File tree

7 files changed

+571
-32
lines changed

7 files changed

+571
-32
lines changed

‎src/main/java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import graphql.schema.DataFetchingEnvironment;
1414
import org.dataloader.DataLoader;
1515
import org.dataloader.DataLoaderRegistry;
16+
import graphql.VisibleForTesting;
1617
import org.jspecify.annotations.NullMarked;
1718
import org.jspecify.annotations.Nullable;
1819

@@ -30,7 +31,8 @@
3031
@NullMarked
3132
public class PerLevelDataLoaderDispatchStrategy implements DataLoaderDispatchStrategy {
3233

33-
private final CallStack initialCallStack;
34+
@VisibleForTesting
35+
final CallStack initialCallStack;
3436
private final ExecutionContext executionContext;
3537
private final boolean enableDataLoaderChaining;
3638

@@ -145,7 +147,8 @@ public void clear() {
145147

146148
}
147149

148-
private static class CallStack {
150+
// package-private for testing
151+
static class CallStack {
149152

150153
/**
151154
* We track three things per level:
@@ -177,8 +180,10 @@ private static class CallStack {
177180
*/
178181

179182
static class StateForLevel {
180-
private final int happenedCompletionFinishedCount;
181-
private final int happenedExecuteObjectCalls;
183+
@VisibleForTesting
184+
final int happenedCompletionFinishedCount;
185+
@VisibleForTesting
186+
final int happenedExecuteObjectCalls;
182187

183188

184189
public StateForLevel() {
@@ -216,7 +221,8 @@ public StateForLevel increaseHappenedExecuteObjectCalls() {
216221

217222
private final Map<Integer, AtomicReference<StateForLevel>> stateForLevelMap = new ConcurrentHashMap<>();
218223

219-
private final Set<Integer> dispatchedLevels = ConcurrentHashMap.newKeySet();
224+
@VisibleForTesting
225+
final Set<Integer> dispatchedLevels = ConcurrentHashMap.newKeySet();
220226

221227
public ChainedDLStack chainedDLStack = new ChainedDLStack();
222228

@@ -439,7 +445,8 @@ private CallStack getCallStack(@Nullable AlternativeCallContext alternativeCallC
439445
}
440446

441447

442-
private boolean markLevelAsDispatchedIfReady(int level, CallStack callStack) {
448+
@VisibleForTesting
449+
boolean markLevelAsDispatchedIfReady(int level, CallStack callStack) {
443450
boolean ready = isLevelReady(level, callStack);
444451
if (ready) {
445452
if (!callStack.dispatchedLevels.add(level)) {

‎src/test/groovy/graphql/execution/instrumentation/dataloader/BatchCompareDataFetchers.java‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616
import java.util.Optional;
1717
import java.util.Random;
1818
import java.util.concurrent.CompletableFuture;
19+
import java.util.concurrent.CountDownLatch;
20+
import java.util.concurrent.ExecutorService;
21+
import java.util.concurrent.Executors;
1922
import java.util.concurrent.atomic.AtomicBoolean;
2023
import java.util.concurrent.atomic.AtomicLong;
2124
import java.util.function.Supplier;
@@ -32,10 +35,21 @@ public class BatchCompareDataFetchers {
3235

3336
AtomicBoolean useAsyncBatchLoading = new AtomicBoolean(false);
3437

38+
private volatile CountDownLatch rootFetcherRendezvous;
39+
private volatile CountDownLatch completionOverlapLatch;
40+
private final AtomicBoolean shopsOverlapSignaled = new AtomicBoolean(false);
41+
private final AtomicBoolean exShopsOverlapSignaled = new AtomicBoolean(false);
42+
private final ExecutorService executor = Executors.newFixedThreadPool(4);
43+
3544
public void useAsyncBatchLoading(boolean flag) {
3645
useAsyncBatchLoading.set(flag);
3746
}
3847

48+
public void useSynchronizedFetching(int numberOfRootFetchers) {
49+
rootFetcherRendezvous = new CountDownLatch(numberOfRootFetchers);
50+
completionOverlapLatch = new CountDownLatch(numberOfRootFetchers);
51+
}
52+
3953

4054
private static final Map<String, Shop> shops = new LinkedHashMap<>();
4155
private static final Map<String, Shop> expensiveShops = new LinkedHashMap<>();
@@ -52,10 +66,10 @@ public void useAsyncBatchLoading(boolean flag) {
5266

5367

5468
public DataFetcher<CompletableFuture<List<Shop>>> shopsDataFetcher =
55-
environment -> supplyAsyncWithSleep(() -> new ArrayList<>(shops.values()));
69+
environment -> supplyAsyncWithRendezvous(() -> new ArrayList<>(shops.values()));
5670

5771
public DataFetcher<CompletableFuture<List<Shop>>> expensiveShopsDataFetcher = environment ->
58-
supplyAsyncWithSleep(() -> new ArrayList<>(expensiveShops.values()));
72+
supplyAsyncWithRendezvous(() -> new ArrayList<>(expensiveShops.values()));
5973

6074
// Departments
6175
private static Map<String, Department> departments = new LinkedHashMap<>();
@@ -101,6 +115,21 @@ private static List<List<Department>> getDepartmentsForShops(List<Shop> shops) {
101115

102116
public DataFetcher<CompletableFuture<List<Department>>> departmentsForShopDataLoaderDataFetcher = environment -> {
103117
Shop shop = environment.getSource();
118+
// When synchronized fetching is enabled, ensure both root fields (shops and expensiveShops)
119+
// are inside their startComplete/stopComplete window before either proceeds.
120+
// This guarantees objectRunningCount never drops to 0 prematurely.
121+
CountDownLatch overlapLatch = completionOverlapLatch;
122+
if (overlapLatch != null) {
123+
AtomicBoolean flag = shop.getId().startsWith("ex") ? exShopsOverlapSignaled : shopsOverlapSignaled;
124+
if (flag.compareAndSet(false, true)) {
125+
overlapLatch.countDown();
126+
try {
127+
overlapLatch.await();
128+
} catch (InterruptedException e) {
129+
throw new RuntimeException(e);
130+
}
131+
}
132+
}
104133
return (CompletableFuture) environment.getDataLoader("departments").load(shop.getId());
105134
};
106135

@@ -149,6 +178,22 @@ private <T> CompletableFuture<T> maybeAsyncWithSleep(Supplier<CompletableFuture<
149178
}
150179
}
151180

181+
private <T> CompletableFuture<T> supplyAsyncWithRendezvous(Supplier<T> supplier) {
182+
CountDownLatch latch = rootFetcherRendezvous;
183+
if (latch != null) {
184+
return CompletableFuture.supplyAsync(() -> {
185+
try {
186+
latch.countDown();
187+
latch.await();
188+
} catch (InterruptedException e) {
189+
throw new RuntimeException(e);
190+
}
191+
return supplier.get();
192+
}, executor);
193+
}
194+
return supplyAsyncWithSleep(supplier);
195+
}
196+
152197
private static <T> CompletableFuture<T> supplyAsyncWithSleep(Supplier<T> supplier) {
153198
Supplier<T> sleepSome = sleepSome(supplier);
154199
return CompletableFuture.supplyAsync(sleepSome);

‎src/test/groovy/graphql/execution/instrumentation/dataloader/DataLoaderPerformanceTest.groovy‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ class DataLoaderPerformanceTest extends Specification {
6060

6161
when:
6262

63+
batchCompareDataFetchers.useSynchronizedFetching(2)
64+
6365
ExecutionInput executionInput = ExecutionInput.newExecutionInput()
6466
.query(getExpensiveQuery(false))
6567
.dataLoaderRegistry(dataLoaderRegistry)
@@ -71,8 +73,8 @@ class DataLoaderPerformanceTest extends Specification {
7173
then:
7274
result.data == expectedExpensiveData
7375

74-
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() <= 2
75-
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() <= 2
76+
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() == 1
77+
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() == 1
7678

7779
where:
7880
incrementalSupport | contextKey
@@ -123,6 +125,7 @@ class DataLoaderPerformanceTest extends Specification {
123125
when:
124126

125127
batchCompareDataFetchers.useAsyncBatchLoading(true)
128+
batchCompareDataFetchers.useSynchronizedFetching(2)
126129

127130
ExecutionInput executionInput = ExecutionInput.newExecutionInput()
128131
.query(getExpensiveQuery(false))
@@ -136,8 +139,8 @@ class DataLoaderPerformanceTest extends Specification {
136139
then:
137140
result.data == expectedExpensiveData
138141

139-
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() <= 2
140-
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() <= 2
142+
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() == 1
143+
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() == 1
141144

142145
where:
143146
incrementalSupport | contextKey

‎src/test/groovy/graphql/execution/instrumentation/dataloader/DataLoaderPerformanceWithChainedInstrumentationTest.groovy‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package graphql.execution.instrumentation.dataloader
33
import graphql.ExecutionInput
44
import graphql.GraphQL
55
import org.dataloader.DataLoaderRegistry
6-
import spock.lang.Ignore
76
import spock.lang.Specification
87

98
import static graphql.ExperimentalApi.ENABLE_INCREMENTAL_SUPPORT
@@ -49,11 +48,12 @@ class DataLoaderPerformanceWithChainedInstrumentationTest extends Specification
4948
incrementalSupport << [true, false]
5049
}
5150

52-
@Ignore("This test flakes on Travis for some reason. Clearly this indicates some sort of problem to investigate. However it also stop releases.")
5351
def "chainedInstrumentation: 970 ensure data loader is performant for multiple field with lists"() {
5452

5553
when:
5654

55+
batchCompareDataFetchers.useSynchronizedFetching(2)
56+
5757
ExecutionInput executionInput = ExecutionInput.newExecutionInput()
5858
.query(getExpensiveQuery(false))
5959
.dataLoaderRegistry(dataLoaderRegistry)
@@ -101,6 +101,7 @@ class DataLoaderPerformanceWithChainedInstrumentationTest extends Specification
101101
when:
102102

103103
batchCompareDataFetchers.useAsyncBatchLoading(true)
104+
batchCompareDataFetchers.useSynchronizedFetching(2)
104105

105106
ExecutionInput executionInput = ExecutionInput.newExecutionInput()
106107
.query(getExpensiveQuery(false))
@@ -112,8 +113,8 @@ class DataLoaderPerformanceWithChainedInstrumentationTest extends Specification
112113
then:
113114
result.data == expectedExpensiveData
114115

115-
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() <= 2
116-
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() <= 2
116+
batchCompareDataFetchers.departmentsForShopsBatchLoaderCounter.get() == 1
117+
batchCompareDataFetchers.productsForDepartmentsBatchLoaderCounter.get() == 1
117118

118119
where:
119120
incrementalSupport << [true, false]

0 commit comments

Comments
 (0)