Skip to content

Commit 4655c4a

Browse files
andimarekandimarek
committed
wip
1 parent aec85bd commit 4655c4a

File tree

1 file changed

+88
-15
lines changed

1 file changed

+88
-15
lines changed

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

Lines changed: 88 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,57 +2,122 @@
22

33
import graphql.agent.result.ExecutionTrackingResult;
44
import graphql.execution.ExecutionContext;
5+
import graphql.execution.ExecutionId;
56
import graphql.execution.ExecutionStrategyParameters;
67
import graphql.execution.ResultPath;
8+
import graphql.schema.DataFetchingEnvironment;
79
import net.bytebuddy.agent.builder.AgentBuilder;
810
import net.bytebuddy.asm.Advice;
11+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
12+
import org.dataloader.DataLoader;
913

1014
import java.lang.instrument.Instrumentation;
15+
import java.util.Map;
16+
import java.util.Set;
1117
import java.util.concurrent.CompletableFuture;
18+
import java.util.concurrent.ConcurrentHashMap;
1219

1320
import static graphql.agent.result.ExecutionTrackingResult.EXECUTION_TRACKING_KEY;
1421
import static net.bytebuddy.matcher.ElementMatchers.nameMatches;
1522
import static net.bytebuddy.matcher.ElementMatchers.named;
1623

1724
public class GraphQLJavaAgent {
1825

26+
public static class ExecutionData {
27+
public Set<ResultPath> resultPaths = ConcurrentHashMap.newKeySet();
28+
29+
public Map<ResultPath, DataLoader> resultPathToDataLoader = new ConcurrentHashMap<>();
30+
public Map<DataLoader, String> dataLoaderToName = new ConcurrentHashMap<>();
31+
32+
@Override
33+
public String toString() {
34+
return "ExecutionData{" +
35+
"resultPathToDataLoader=" + resultPathToDataLoader +
36+
", dataLoaderToName=" + dataLoaderToName +
37+
'}';
38+
}
39+
}
40+
41+
public static final Map<ExecutionId, ExecutionData> executionDataMap = new ConcurrentHashMap<>();
42+
public static final Map<DataLoader, ExecutionId> dataLoaderToExecutionId = new ConcurrentHashMap<>();
43+
1944
public static void agentmain(String agentArgs, Instrumentation inst) {
2045
System.out.println("Agent is running");
2146
new AgentBuilder.Default()
22-
.disableClassFormatChanges()
47+
.type(named("graphql.execution.Execution"))
48+
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> {
49+
System.out.println("transforming " + typeDescription);
50+
return builder
51+
.visit(Advice.to(ExecutionAdvice.class).on(nameMatches("executeOperation")));
52+
53+
})
2354
.type(named("graphql.execution.ExecutionStrategy"))
2455
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> {
2556
System.out.println("Transforming " + typeDescription);
2657
return builder
2758
.visit(Advice.to(DataFetcherInvokeAdvice.class).on(nameMatches("invokeDataFetcher")));
2859
})
29-
.installOn(inst);
30-
31-
new AgentBuilder.Default()
32-
.disableClassFormatChanges()
33-
.type(named("graphql.execution.Execution"))
60+
.type(named("org.dataloader.DataLoaderRegistry"))
3461
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> {
3562
System.out.println("transforming " + typeDescription);
3663
return builder
37-
.visit(Advice.to(ExecutionAdvice.class).on(nameMatches("executeOperation")));
38-
}).installOn(inst);
39-
40-
new AgentBuilder.Default()
41-
.disableClassFormatChanges()
42-
.with(AgentBuilder.RedefinitionStrategy.DISABLED)
43-
.type(named("org.dataloader.DataLoaderRegistry"))
64+
.visit(Advice.to(DataLoaderRegistryAdvice.class).on(nameMatches("dispatchAll")));
65+
})
66+
.type(named("org.dataloader.DataLoader"))
4467
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> {
4568
System.out.println("transforming " + typeDescription);
4669
return builder
47-
.visit(Advice.to(DataLoaderRegistryAdvice.class).on(nameMatches("dispatchAll")));
48-
}).installOn(inst);
70+
.visit(Advice.to(DataLoaderAdvice.class).on(nameMatches("load")));
71+
})
72+
.type(named("graphql.schema.DataFetchingEnvironmentImpl"))
73+
.transform((builder, typeDescription, classLoader, module, protectionDomain) -> {
74+
System.out.println("transforming " + typeDescription);
75+
return builder
76+
.visit(Advice.to(DataFetchingEnvironmentAdvice.class).on(nameMatches("getDataLoader")));
77+
})
78+
.installOn(inst);
79+
4980
}
5081
}
5182

83+
class DataFetchingEnvironmentAdvice {
84+
85+
86+
@Advice.OnMethodExit
87+
public static void getDataLoader(@Advice.Argument(0) String dataLoaderName,
88+
@Advice.This(typing = Assigner.Typing.DYNAMIC) DataFetchingEnvironment dataFetchingEnvironment,
89+
@Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) DataLoader dataLoader) {
90+
GraphQLJavaAgent.ExecutionData executionData = GraphQLJavaAgent.executionDataMap.get(dataFetchingEnvironment.getExecutionId());
91+
System.out.println("execution data: " + executionData);
92+
executionData.resultPathToDataLoader.put(dataFetchingEnvironment.getExecutionStepInfo().getPath(), dataLoader);
93+
executionData.dataLoaderToName.put(dataLoader, dataLoaderName);
94+
GraphQLJavaAgent.dataLoaderToExecutionId.put(dataLoader, dataFetchingEnvironment.getExecutionId());
95+
96+
System.out.println(dataLoaderName + " > " + dataLoader);
97+
}
98+
99+
}
100+
101+
102+
class DataLoaderAdvice {
103+
104+
@Advice.OnMethodEnter
105+
public static void load(@Advice.This(typing = Assigner.Typing.DYNAMIC) Object dataLoader) {
106+
ExecutionId executionId = GraphQLJavaAgent.dataLoaderToExecutionId.get(dataLoader);
107+
String dataLoaderName = GraphQLJavaAgent.executionDataMap.get(executionId).dataLoaderToName.get(dataLoader);
108+
System.out.println("dataloader " + dataLoaderName + " load for execution " + executionId);
109+
}
110+
111+
}
112+
52113
class DataLoaderRegistryAdvice {
53114

54115
@Advice.OnMethodEnter
55116
public static void dispatchAll() {
117+
// StackWalker instance = StackWalker.getInstance();
118+
// instance.forEach(stackFrame -> {
119+
// System.out.println(stackFrame.getClassName() + " " + stackFrame.getMethodName() + " " + stackFrame.getLineNumber());
120+
// });
56121
System.out.println("calling dispatchAll");
57122
}
58123

@@ -62,14 +127,22 @@ class ExecutionAdvice {
62127

63128
@Advice.OnMethodEnter
64129
public static void executeOperationEnter(@Advice.Argument(0) ExecutionContext executionContext) {
130+
System.out.println("execution started for: " + executionContext.getExecutionId());
65131
executionContext.getGraphQLContext().put(EXECUTION_TRACKING_KEY, new ExecutionTrackingResult());
132+
GraphQLJavaAgent.executionDataMap.put(executionContext.getExecutionId(), new GraphQLJavaAgent.ExecutionData());
133+
}
134+
135+
@Advice.OnMethodExit
136+
public static void executeOperationExit(@Advice.Argument(0) ExecutionContext executionContext) {
137+
System.out.println("execution finished for: " + executionContext.getExecutionId());
66138
}
67139
}
68140

69141
class DataFetcherInvokeAdvice {
70142
@Advice.OnMethodEnter
71143
public static void invokeDataFetcherEnter(@Advice.Argument(0) ExecutionContext executionContext,
72144
@Advice.Argument(1) ExecutionStrategyParameters parameters) {
145+
73146
ExecutionTrackingResult executionTrackingResult = executionContext.getGraphQLContext().get(EXECUTION_TRACKING_KEY);
74147
executionTrackingResult.start(parameters.getPath(), System.nanoTime());
75148
}

0 commit comments

Comments
 (0)