Skip to content

Commit 7a1aa1c

Browse files
evanchoolyevanchooly
committed
POC of cleaning up the initial code origin set up instrumentation
1 parent b00c24e commit 7a1aa1c

File tree

35 files changed

+236
-159
lines changed

35 files changed

+236
-159
lines changed

‎dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java

+5-11
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ public interface ExceptionDebugger {
116116
public interface CodeOriginRecorder {
117117
String captureCodeOrigin(boolean entry);
118118

119-
String captureCodeOrigin(Method method, boolean entry, boolean instrument);
119+
String captureCodeOrigin(Method method, boolean entry);
120120
}
121121

122122
private static volatile ProductConfigUpdater productConfigUpdater;
@@ -466,32 +466,26 @@ public static void commit(
466466
}
467467
}
468468

469-
public static String captureCodeOrigin(boolean entry) {
469+
public static void captureCodeOrigin(boolean entry) {
470470
try {
471471
CodeOriginRecorder recorder = codeOriginRecorder;
472472
if (recorder != null) {
473-
return recorder.captureCodeOrigin(entry);
473+
recorder.captureCodeOrigin(entry);
474474
}
475475
} catch (Exception ex) {
476476
LOGGER.debug("Error in captureCodeOrigin: ", ex);
477477
}
478-
return null;
479478
}
480479

481-
public static String captureCodeOrigin(Method method, boolean entry) {
482-
return captureCodeOrigin(method, entry, true);
483-
}
484-
485-
public static String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
480+
public static void captureCodeOrigin(Method method, boolean entry) {
486481
try {
487482
CodeOriginRecorder recorder = codeOriginRecorder;
488483
if (recorder != null) {
489-
return recorder.captureCodeOrigin(method, entry, instrument);
484+
recorder.captureCodeOrigin(method, entry);
490485
}
491486
} catch (Exception ex) {
492487
LOGGER.debug("Error in captureCodeOrigin: ", ex);
493488
}
494-
return null;
495489
}
496490

497491
public static void handleException(Throwable t, AgentSpan span) {

‎dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/spanorigin/CodeOriginInfo.java

-24
This file was deleted.

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java

+18-25
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
4545

4646
private final int maxUserFrames;
4747

48-
private AgentTaskScheduler scheduler = AgentTaskScheduler.INSTANCE;
48+
private AgentTaskScheduler scheduler;
4949

5050
public DefaultCodeOriginRecorder(Config config, ConfigurationUpdater configurationUpdater) {
5151
this(config, configurationUpdater, AgentTaskScheduler.INSTANCE);
@@ -70,51 +70,44 @@ public String captureCodeOrigin(boolean entry) {
7070
element.getMethodName(),
7171
null,
7272
String.valueOf(element.getLineNumber()));
73-
probe = createProbe(fingerprint, entry, where, true);
73+
probe = createProbe(fingerprint, entry, where);
7474

7575
LOG.debug("Creating probe for location {}", where);
7676
}
7777
return probe.getId();
7878
}
7979

8080
@Override
81-
public String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
81+
public String captureCodeOrigin(Method method, boolean entry) {
8282
String fingerprint = method.toString();
8383
CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
8484
if (probe == null) {
85-
probe = createProbe(fingerprint, entry, Where.of(method), instrument);
85+
probe = createProbe(fingerprint, entry, Where.of(method));
8686
LOG.debug("Creating probe for method {}", fingerprint);
87-
} else if (!instrument) {
88-
// direct call to fill code origin info without using probe instrumentation
89-
// buildLocation should be called before in order to gather location info
90-
probe.commit(
91-
CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
9287
}
9388
return probe.getId();
9489
}
9590

9691
public void registerLogProbe(CodeOriginProbe probe) {
97-
LogProbe logProbe =
98-
logProbes.computeIfAbsent(
99-
probe.getId(),
100-
key ->
101-
new Builder()
102-
.language(probe.getLanguage())
103-
.probeId(ProbeId.newId())
104-
.where(probe.getWhere())
105-
.evaluateAt(probe.getEvaluateAt())
106-
.captureSnapshot(true)
107-
.tags("session_id:*")
108-
.snapshotProcessor(new CodeOriginSnapshotConsumer(probe.entrySpanProbe()))
109-
.build());
92+
logProbes.computeIfAbsent(
93+
probe.getId(),
94+
key ->
95+
new Builder()
96+
.language(probe.getLanguage())
97+
.probeId(ProbeId.newId())
98+
.where(probe.getWhere())
99+
.evaluateAt(probe.getEvaluateAt())
100+
.captureSnapshot(true)
101+
.tags("session_id:*")
102+
.snapshotProcessor(new CodeOriginSnapshotConsumer(probe.entrySpanProbe()))
103+
.build());
110104
}
111105

112-
private CodeOriginProbe createProbe(
113-
String fingerPrint, boolean entry, Where where, boolean instrument) {
106+
private CodeOriginProbe createProbe(String fingerPrint, boolean entry, Where where) {
114107
CodeOriginProbe probe;
115108
AgentSpan span = AgentTracer.activeSpan();
116109

117-
probe = new CodeOriginProbe(ProbeId.newId(), entry, where, instrument);
110+
probe = new CodeOriginProbe(ProbeId.newId(), entry, where);
118111
addFingerprint(fingerPrint, probe);
119112
CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);
120113

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/ASMHelper.java

+24-1
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,25 @@
11
package com.datadog.debugger.instrumentation;
22

33
import static com.datadog.debugger.instrumentation.Types.REFLECTIVE_FIELD_VALUE_RESOLVER_TYPE;
4+
import static java.lang.String.format;
45
import static org.objectweb.asm.Type.getMethodDescriptor;
56
import static org.objectweb.asm.Type.getObjectType;
67

78
import com.datadog.debugger.agent.Generated;
89
import datadog.trace.util.Strings;
10+
import de.thetaphi.forbiddenapis.SuppressForbidden;
11+
import java.io.PrintWriter;
12+
import java.io.StringWriter;
913
import java.lang.reflect.Field;
1014
import java.lang.reflect.Modifier;
1115
import java.util.ArrayList;
16+
import java.util.Arrays;
1217
import java.util.Collections;
1318
import java.util.Comparator;
1419
import java.util.List;
20+
import java.util.StringJoiner;
1521
import java.util.stream.Collectors;
1622
import org.objectweb.asm.Opcodes;
17-
import org.objectweb.asm.Type;
1823
import org.objectweb.asm.signature.SignatureReader;
1924
import org.objectweb.asm.signature.SignatureVisitor;
2025
import org.objectweb.asm.tree.AbstractInsnNode;
@@ -28,6 +33,8 @@
2833
import org.objectweb.asm.tree.MethodInsnNode;
2934
import org.objectweb.asm.tree.MethodNode;
3035
import org.objectweb.asm.tree.TypeInsnNode;
36+
import org.objectweb.asm.util.Textifier;
37+
import org.objectweb.asm.util.TraceClassVisitor;
3138

3239
/** Helper class for bytecode generation */
3340
public class ASMHelper {
@@ -36,6 +43,22 @@ public class ASMHelper {
3643
public static final Type STRING_TYPE = new Type(Types.STRING_TYPE);
3744
public static final Type LONG_TYPE = new Type(org.objectweb.asm.Type.LONG_TYPE);
3845

46+
@SuppressForbidden
47+
static String extractMethod(ClassNode classNode, String method) {
48+
StringJoiner joiner = new StringJoiner("\n");
49+
joiner.add("Class: " + classNode.name);
50+
StringWriter writer = new StringWriter();
51+
classNode.accept(new TraceClassVisitor(null, new Textifier(), new PrintWriter(writer)));
52+
List<String> strings = Arrays.asList(writer.toString().split("\n"));
53+
for (int i = 0; i < strings.size(); i++) {
54+
if (strings.get(i).matches(format(".*(private|public).* %s\\(.*", method))) {
55+
while (!strings.get(i).equals(""))
56+
joiner.add(String.format("[%3d] %s", i, strings.get(i++)));
57+
}
58+
}
59+
return joiner.toString();
60+
}
61+
3962
public static void invokeInterface(
4063
InsnList insnList,
4164
org.objectweb.asm.Type owner,

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CodeOriginInstrumentor.java

+57-8
Original file line numberDiff line numberDiff line change
@@ -5,35 +5,84 @@
55
import static com.datadog.debugger.instrumentation.Types.DEBUGGER_CONTEXT_TYPE;
66
import static com.datadog.debugger.instrumentation.Types.STRING_TYPE;
77
import static java.lang.Integer.parseInt;
8+
import static java.lang.String.format;
89

910
import com.datadog.debugger.instrumentation.InstrumentationResult.Status;
1011
import com.datadog.debugger.probe.CodeOriginProbe;
1112
import com.datadog.debugger.probe.ProbeDefinition;
13+
import datadog.trace.bootstrap.debugger.DebuggerContext;
1214
import datadog.trace.bootstrap.debugger.ProbeId;
15+
import java.util.ArrayList;
1316
import java.util.List;
17+
import java.util.ListIterator;
1418
import org.objectweb.asm.Type;
1519
import org.objectweb.asm.tree.AbstractInsnNode;
1620
import org.objectweb.asm.tree.InsnList;
1721
import org.objectweb.asm.tree.LabelNode;
22+
import org.objectweb.asm.tree.MethodInsnNode;
1823

1924
public class CodeOriginInstrumentor extends Instrumentor {
25+
private static final String CAPTURE;
26+
27+
static {
28+
String className = DebuggerContext.class.getName().replace('.', '/');
29+
CAPTURE = format("%s#%s", className, "captureCodeOrigin");
30+
}
31+
32+
private static String buildDescription(AbstractInsnNode node) {
33+
if (!(node instanceof MethodInsnNode)) return "";
34+
MethodInsnNode method = (MethodInsnNode) node;
35+
return format("%s#%s", method.owner, method.name);
36+
}
37+
2038
public CodeOriginInstrumentor(
21-
ProbeDefinition definition,
22-
MethodInfo methodInfo,
23-
List<DiagnosticMessage> diagnostics,
24-
List<ProbeId> probeIds) {
25-
super(definition, methodInfo, diagnostics, probeIds);
39+
ProbeDefinition definition, MethodInfo methodInfo, List<ProbeId> probeIds) {
40+
super(definition, methodInfo, null, probeIds);
2641
}
2742

2843
@Override
2944
public Status instrument() {
30-
InsnList insnList = new InsnList();
45+
AbstractInsnNode insertionPoint = stripSetup();
46+
methodNode.instructions.insert(
47+
insertionPoint != null ? insertionPoint : findInsertionPoint(), codeOriginCall());
3148

49+
return Status.INSTALLED;
50+
}
51+
52+
private InsnList codeOriginCall() {
53+
InsnList insnList = new InsnList();
3254
ldc(insnList, probeIds.get(0).getEncodedId());
3355
invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "codeOrigin", Type.VOID_TYPE, STRING_TYPE);
34-
methodNode.instructions.insert(findInsertionPoint(), insnList);
56+
return insnList;
57+
}
3558

36-
return InstrumentationResult.Status.INSTALLED;
59+
private AbstractInsnNode stripSetup() {
60+
ListIterator<AbstractInsnNode> iterator = methodNode.instructions.iterator();
61+
List<AbstractInsnNode> list = new ArrayList<>();
62+
AbstractInsnNode insertionPoint = null;
63+
while (iterator.hasNext()) {
64+
AbstractInsnNode next = iterator.next();
65+
if (buildDescription(next).equals(CAPTURE)) {
66+
int argumentCount = Type.getMethodType(((MethodInsnNode) next).desc).getArgumentCount();
67+
int count;
68+
if (argumentCount == 2) {
69+
count = 10;
70+
} else if (argumentCount == 1) {
71+
count = 1;
72+
} else {
73+
return null;
74+
}
75+
for (int i = 0; i < count; i++) {
76+
list.remove(list.size() - 1);
77+
}
78+
insertionPoint = list.get(list.size() - 1);
79+
} else {
80+
list.add(next);
81+
}
82+
}
83+
methodNode.instructions = new InsnList();
84+
list.forEach(n -> methodNode.instructions.add(n));
85+
return insertionPoint;
3786
}
3887

3988
private AbstractInsnNode findInsertionPoint() {

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java

+2-7
Original file line numberDiff line numberDiff line change
@@ -26,23 +26,18 @@
2626
public class CodeOriginProbe extends ProbeDefinition {
2727
private static final Logger LOGGER = LoggerFactory.getLogger(CodeOriginProbe.class);
2828

29-
private final boolean instrument;
3029
private final boolean entrySpanProbe;
3130
private String signature;
3231

33-
public CodeOriginProbe(ProbeId probeId, boolean entry, Where where, boolean instrument) {
32+
public CodeOriginProbe(ProbeId probeId, boolean entry, Where where) {
3433
super(LANGUAGE, probeId, (Tag[]) null, where, MethodLocation.ENTRY);
35-
this.instrument = instrument;
3634
this.entrySpanProbe = entry;
3735
}
3836

3937
@Override
4038
public Status instrument(
4139
MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
42-
if (instrument) {
43-
return new CodeOriginInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
44-
}
45-
return Status.INSTALLED;
40+
return new CodeOriginInstrumentor(this, methodInfo, probeIds).instrument();
4641
}
4742

4843
@Override

‎dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOrigin.java

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package com.datadog.debugger.origin;
2+
3+
import java.lang.annotation.ElementType;
4+
import java.lang.annotation.Target;
5+
6+
@Target(ElementType.METHOD)
7+
public @interface CodeOrigin {}

0 commit comments

Comments
 (0)