Skip to content

Commit 1b4c2fb

Browse files
andimarekandimarek
committed
wip
1 parent c8b0f31 commit 1b4c2fb

File tree

2 files changed

+49
-62
lines changed

2 files changed

+49
-62
lines changed

‎agent/src/main/java/graphql/agent/AfterExecutionHandler.java‎

Lines changed: 0 additions & 25 deletions
This file was deleted.

‎agent/src/main/java/graphql/agent/GraphQLJavaAgent.java‎

Lines changed: 49 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import java.util.concurrent.ConcurrentHashMap;
2323
import java.util.concurrent.atomic.AtomicLong;
2424
import java.util.concurrent.atomic.AtomicReference;
25+
import java.util.function.BiConsumer;
2526

2627
import static graphql.agent.GraphQLJavaAgent.ExecutionData.DFResultType.DONE_CANCELLED;
2728
import static graphql.agent.GraphQLJavaAgent.ExecutionData.DFResultType.DONE_EXCEPTIONALLY;
@@ -193,6 +194,54 @@ public static void agentmain(String agentArgs, Instrumentation inst) {
193194
.installOn(inst);
194195

195196
}
197+
198+
public static class ExecutionAdvice {
199+
200+
public static class AfterExecutionHandler implements BiConsumer<Object, Throwable> {
201+
202+
private final ExecutionContext executionContext;
203+
204+
public AfterExecutionHandler(ExecutionContext executionContext) {
205+
this.executionContext = executionContext;
206+
}
207+
208+
public void accept(Object o, Throwable throwable) {
209+
ExecutionId executionId = executionContext.getExecutionId();
210+
GraphQLJavaAgent.ExecutionData executionData = GraphQLJavaAgent.executionIdToData.get(executionId);
211+
executionData.endExecutionTime.set(System.nanoTime());
212+
System.out.println("execution finished for: " + executionId + " with data " + executionData);
213+
System.out.println(executionData.print(executionId.toString()));
214+
}
215+
216+
}
217+
218+
219+
@Advice.OnMethodEnter
220+
public static void executeOperationEnter(@Advice.Argument(0) ExecutionContext executionContext) {
221+
GraphQLJavaAgent.ExecutionData executionData = new GraphQLJavaAgent.ExecutionData();
222+
executionData.startExecutionTime.set(System.nanoTime());
223+
executionData.startThread.set(Thread.currentThread().getName());
224+
System.out.println("execution started for: " + executionContext.getExecutionId());
225+
executionContext.getGraphQLContext().put(EXECUTION_TRACKING_KEY, new ExecutionTrackingResult());
226+
227+
GraphQLJavaAgent.executionIdToData.put(executionContext.getExecutionId(), executionData);
228+
229+
DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
230+
for (String name : dataLoaderRegistry.getDataLoadersMap().keySet()) {
231+
DataLoader dataLoader = dataLoaderRegistry.getDataLoader(name);
232+
GraphQLJavaAgent.dataLoaderToExecutionId.put(dataLoader, executionContext.getExecutionId());
233+
executionData.dataLoaderToName.put(dataLoader, name);
234+
}
235+
}
236+
237+
@Advice.OnMethodExit
238+
public static void executeOperationExit(@Advice.Argument(0) ExecutionContext executionContext,
239+
@Advice.Return(typing = Assigner.Typing.DYNAMIC) CompletableFuture<Object> result) {
240+
241+
result.whenComplete(new AfterExecutionHandler(executionContext));
242+
}
243+
}
244+
196245
}
197246

198247
class DataFetchingEnvironmentAdvice {
@@ -260,43 +309,6 @@ public static void dispatchAll(@Advice.This(typing = Assigner.Typing.DYNAMIC) Ob
260309

261310
}
262311

263-
class ExecutionAdvice{
264-
265-
@Advice.OnMethodEnter
266-
public static void executeOperationEnter(@Advice.Argument(0) ExecutionContext executionContext) {
267-
GraphQLJavaAgent.ExecutionData executionData = new GraphQLJavaAgent.ExecutionData();
268-
executionData.startExecutionTime.set(System.nanoTime());
269-
executionData.startThread.set(Thread.currentThread().getName());
270-
System.out.println("execution started for: " + executionContext.getExecutionId());
271-
executionContext.getGraphQLContext().put(EXECUTION_TRACKING_KEY, new ExecutionTrackingResult());
272-
273-
GraphQLJavaAgent.executionIdToData.put(executionContext.getExecutionId(), executionData);
274-
275-
DataLoaderRegistry dataLoaderRegistry = executionContext.getDataLoaderRegistry();
276-
for (String name : dataLoaderRegistry.getDataLoadersMap().keySet()) {
277-
DataLoader dataLoader = dataLoaderRegistry.getDataLoader(name);
278-
GraphQLJavaAgent.dataLoaderToExecutionId.put(dataLoader, executionContext.getExecutionId());
279-
executionData.dataLoaderToName.put(dataLoader, name);
280-
}
281-
}
282-
283-
@Advice.OnMethodExit
284-
public static void executeOperationExit(@Advice.Argument(0) ExecutionContext executionContext,
285-
@Advice.Return(typing = Assigner.Typing.DYNAMIC) CompletableFuture<Object> result) {
286-
287-
result.whenComplete(new AfterExecutionHandler(executionContext));
288-
// ExecutionId executionId = executionContext.getExecutionId();
289-
// GraphQLJavaAgent.ExecutionData executionData = GraphQLJavaAgent.executionIdToData.get(executionId);
290-
// executionData.endExecutionTime.set(System.nanoTime());
291-
// System.out.println("execution finished for: " + executionId + " with data " + executionData);
292-
// System.out.println(executionData.print(executionId.toString()));
293-
// cleanup
294-
// GraphQLJavaAgent.executionDataMap.get(executionId).dataLoaderToName.forEach((dataLoader, s) -> {
295-
// GraphQLJavaAgent.dataLoaderToExecutionId.remove(dataLoader);
296-
// });
297-
// GraphQLJavaAgent.executionDataMap.remove(executionContext.getExecutionId());
298-
}
299-
}
300312

301313
class DataFetcherInvokeAdvice {
302314
@Advice.OnMethodEnter

0 commit comments

Comments
 (0)