Skip to content

Commit e6faa18

Browse files
andimarekandimarek
committed
manual dispatch collected
1 parent 0b7c15b commit e6faa18

File tree

5 files changed

+99
-9
lines changed

5 files changed

+99
-9
lines changed

‎src/main/java/graphql/Profiler.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,11 @@ default void batchLoadedOldStrategy(String name, int level, int count) {
5151

5252
}
5353

54-
default void batchLoadedNewStrategy(String name, @Nullable Integer level, int count) {
54+
default void batchLoadedNewStrategy(String dataLoaderName, @Nullable Integer level, int count) {
55+
56+
}
57+
58+
default <V> void manualDispatch(String dataLoaderName, int level, int count) {
5559

5660
}
5761
}

‎src/main/java/graphql/ProfilerImpl.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -138,11 +138,16 @@ public void oldStrategyDispatchingAll(int level) {
138138

139139
@Override
140140
public void batchLoadedOldStrategy(String name, int level, int count) {
141-
profilerResult.addDispatchEvent(name, level, count);
141+
profilerResult.addDispatchEvent(name, level, count, ProfilerResult.DispatchEventType.STRATEGY_DISPATCH);
142142
}
143143

144144
@Override
145-
public void batchLoadedNewStrategy(String name, @Nullable Integer level, int count) {
146-
profilerResult.addDispatchEvent(name, level, count);
145+
public void batchLoadedNewStrategy(String dataLoaderName, @Nullable Integer level, int count) {
146+
profilerResult.addDispatchEvent(dataLoaderName, level, count, ProfilerResult.DispatchEventType.STRATEGY_DISPATCH);
147+
}
148+
149+
@Override
150+
public <V> void manualDispatch(String dataLoaderName, int level, int count) {
151+
profilerResult.addDispatchEvent(dataLoaderName, level, count, ProfilerResult.DispatchEventType.MANUAL_DISPATCH);
147152
}
148153
}

‎src/main/java/graphql/ProfilerResult.java‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -63,15 +63,22 @@ public void setInstrumentationClasses(List<String> instrumentationClasses) {
6363
}
6464

6565

66+
public enum DispatchEventType {
67+
STRATEGY_DISPATCH,
68+
MANUAL_DISPATCH,
69+
}
70+
6671
public static class DispatchEvent {
6772
final String dataLoaderName;
6873
final @Nullable Integer level; // is null for delayed dispatching
6974
final int count; // how many
75+
final DispatchEventType type;
7076

71-
public DispatchEvent(String dataLoaderName, @Nullable Integer level, int count) {
77+
public DispatchEvent(String dataLoaderName, @Nullable Integer level, int count, DispatchEventType type) {
7278
this.dataLoaderName = dataLoaderName;
7379
this.level = level;
7480
this.count = count;
81+
this.type = type;
7582
}
7683

7784
public String getDataLoaderName() {
@@ -86,10 +93,15 @@ public int getCount() {
8693
return count;
8794
}
8895

96+
public DispatchEventType getType() {
97+
return type;
98+
}
99+
89100
@Override
90101
public String toString() {
91102
return "DispatchEvent{" +
92-
"dataLoaderName='" + dataLoaderName + '\'' +
103+
"type=" + type +
104+
", dataLoaderName='" + dataLoaderName + '\'' +
93105
", level=" + level +
94106
", count=" + count +
95107
'}';
@@ -164,8 +176,8 @@ void chainedStrategyDispatching(int level) {
164176
chainedStrategyDispatching.add(level);
165177
}
166178

167-
void addDispatchEvent(String dataLoaderName, @Nullable Integer level, int count) {
168-
dispatchEvents.add(new DispatchEvent(dataLoaderName, level, count));
179+
void addDispatchEvent(String dataLoaderName, @Nullable Integer level, int count, DispatchEventType type) {
180+
dispatchEvents.add(new DispatchEvent(dataLoaderName, level, count, type));
169181
}
170182

171183
// public getters
@@ -358,6 +370,7 @@ public List<Map<String, Object>> getDispatchEventsAsMap() {
358370

359371
@Override
360372
public String toString() {
361-
return shortSummary();
373+
return "ProfilerResult" + shortSummaryMap();
362374
}
375+
363376
}

‎src/main/java/graphql/schema/DataLoaderWithContext.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import org.jspecify.annotations.NullMarked;
1010
import org.jspecify.annotations.Nullable;
1111

12+
import java.util.List;
1213
import java.util.concurrent.CompletableFuture;
1314

1415
@Internal
@@ -40,4 +41,15 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
4041
return result;
4142
}
4243

44+
@Override
45+
public CompletableFuture<List<V>> dispatch() {
46+
CompletableFuture<List<V>> dispatchResult = delegate.dispatch();
47+
dispatchResult.whenComplete((result, error) -> {
48+
if (result != null) {
49+
DataFetchingEnvironmentImpl.DFEInternalState dfeInternalState = (DataFetchingEnvironmentImpl.DFEInternalState) dfe.toInternal();
50+
dfeInternalState.getProfiler().manualDispatch(dataLoaderName, dfe.getExecutionStepInfo().getPath().getLevel(), result.size());
51+
}
52+
});
53+
return dispatchResult;
54+
}
4355
}

‎src/test/groovy/graphql/ProfilerTest.groovy‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import spock.lang.Specification
1515

1616
import java.time.Duration
1717
import java.util.concurrent.CompletableFuture
18+
import java.util.concurrent.atomic.AtomicInteger
1819

1920
import static graphql.ExecutionInput.newExecutionInput
2021
import static graphql.ProfilerResult.DataFetcherResultType.COMPLETABLE_FUTURE_COMPLETED
@@ -54,6 +55,61 @@ class ProfilerTest extends Specification {
5455

5556
}
5657

58+
def "manual dataloader dispatch"() {
59+
given:
60+
def sdl = '''
61+
62+
type Query {
63+
dog: String
64+
}
65+
'''
66+
AtomicInteger batchLoadCalls = new AtomicInteger()
67+
BatchLoader<String, String> batchLoader = { keys ->
68+
return supplyAsync {
69+
batchLoadCalls.incrementAndGet()
70+
Thread.sleep(250)
71+
println "BatchLoader called with keys: $keys"
72+
return ["Luna"]
73+
}
74+
}
75+
76+
DataLoader<String, String> nameDataLoader = DataLoaderFactory.newDataLoader(batchLoader);
77+
78+
DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry();
79+
dataLoaderRegistry.register("name", nameDataLoader);
80+
81+
def df1 = { env ->
82+
def loader = env.getDataLoader("name")
83+
def result = loader.load("Key1")
84+
loader.dispatch()
85+
return result
86+
} as DataFetcher
87+
88+
def fetchers = ["Query": ["dog": df1]]
89+
def schema = TestUtil.schema(sdl, fetchers)
90+
def graphQL = GraphQL.newGraphQL(schema).build()
91+
92+
def query = "{ dog } "
93+
def ei = newExecutionInput(query).dataLoaderRegistry(dataLoaderRegistry).profileExecution(true).build()
94+
setEnableDataLoaderChaining(ei.graphQLContext, true)
95+
96+
when:
97+
def efCF = graphQL.executeAsync(ei)
98+
Awaitility.await().until { efCF.isDone() }
99+
def er = efCF.get()
100+
def profilerResult = ei.getGraphQLContext().get(ProfilerResult.PROFILER_CONTEXT_KEY) as ProfilerResult
101+
then:
102+
er.data == [dog: "Luna"]
103+
batchLoadCalls.get() == 1
104+
then:
105+
profilerResult.getDispatchEvents()[0].type == ProfilerResult.DispatchEventType.MANUAL_DISPATCH
106+
profilerResult.getDispatchEvents()[0].dataLoaderName == "name"
107+
profilerResult.getDispatchEvents()[0].count == 1
108+
profilerResult.getDispatchEvents()[0].level == 1
109+
110+
}
111+
112+
57113
def "collects instrumentation list"() {
58114
given:
59115
def sdl = '''

0 commit comments

Comments
 (0)