Skip to content

Commit aff1e5e

Browse files
jpbempeljpbempel
committed
Add support of Kotlin's SourceDebugExtension
In Kotlin some constructions remap source line beyond of the end of file (ex: stream operations, ...) which leads to incorrect line numbering in SymDB. The Kotlin classfiles contains information about this line remapping using the SourceDebugExtension (JSR-45). We are leveraging this extension to remap correctly the line numbers and send corrected lines to SymDB. We are introducing a SourceRemapper interface with one implementation for Kotlin for now, but keeping open door for other languages like Scala.
1 parent 5a54a81 commit aff1e5e

File tree

11 files changed

+379
-136
lines changed

11 files changed

+379
-136
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import com.datadog.debugger.probe.Where;
3131
import com.datadog.debugger.sink.Snapshot;
3232
import com.datadog.debugger.util.ClassFileLines;
33+
import com.datadog.debugger.util.JvmLanguage;
3334
import datadog.trace.api.Config;
3435
import datadog.trace.bootstrap.debugger.Limits;
3536
import datadog.trace.bootstrap.debugger.MethodLocation;

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

Lines changed: 2 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.datadog.debugger.instrumentation.DiagnosticMessage.Kind;
1010
import com.datadog.debugger.probe.ProbeDefinition;
1111
import com.datadog.debugger.util.ClassFileLines;
12+
import com.datadog.debugger.util.JvmLanguage;
1213
import java.util.ArrayList;
1314
import java.util.Arrays;
1415
import java.util.HashMap;
@@ -69,7 +70,7 @@ public Instrumenter(
6970
argOffset += t.getSize();
7071
}
7172
localVarsBySlotArray = extractLocalVariables(argTypes);
72-
this.language = JvmLanguage.of(classNode);
73+
this.language = JvmLanguage.of(classNode.sourceFile);
7374
}
7475

7576
public abstract InstrumentationResult.Status instrument();
@@ -295,31 +296,4 @@ public FinallyBlock(LabelNode startLabel, LabelNode endLabel, LabelNode handlerL
295296
this.handlerLabel = handlerLabel;
296297
}
297298
}
298-
299-
protected enum JvmLanguage {
300-
JAVA,
301-
KOTLIN,
302-
SCALA,
303-
GROOVY,
304-
UNKNOWN;
305-
306-
public static JvmLanguage of(ClassNode classNode) {
307-
if (classNode.sourceFile == null) {
308-
return UNKNOWN;
309-
}
310-
if (classNode.sourceFile.endsWith(".java")) {
311-
return JAVA;
312-
}
313-
if (classNode.sourceFile.endsWith(".kt")) {
314-
return KOTLIN;
315-
}
316-
if (classNode.sourceFile.endsWith(".scala")) {
317-
return SCALA;
318-
}
319-
if (classNode.sourceFile.endsWith(".groovy")) {
320-
return GROOVY;
321-
}
322-
return UNKNOWN;
323-
}
324-
}
325299
}

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SourceRemapper.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package com.datadog.debugger.symbol;
2+
3+
import com.datadog.debugger.util.JvmLanguage;
4+
import datadog.trace.agent.tooling.stratum.SourceMap;
5+
import datadog.trace.agent.tooling.stratum.StratumExt;
6+
import datadog.trace.api.Pair;
7+
8+
public interface SourceRemapper {
9+
10+
int remapSourceLine(int line);
11+
12+
static SourceRemapper getSourceRemapper(String sourceFile, SourceMap sourceMap) {
13+
JvmLanguage jvmLanguage = JvmLanguage.of(sourceFile);
14+
switch (jvmLanguage) {
15+
case KOTLIN:
16+
StratumExt stratum = sourceMap.getStratum("KotlinDebug");
17+
if (stratum == null) {
18+
throw new IllegalArgumentException("No stratum found for KotlinDebug");
19+
}
20+
return new KotlinSourceRemapper(stratum);
21+
default:
22+
return NOOP_REMAPPER;
23+
}
24+
}
25+
26+
SourceRemapper NOOP_REMAPPER = new NoopSourceRemapper();
27+
28+
class NoopSourceRemapper implements SourceRemapper {
29+
@Override
30+
public int remapSourceLine(int line) {
31+
return line;
32+
}
33+
}
34+
35+
class KotlinSourceRemapper implements SourceRemapper {
36+
private final StratumExt stratum;
37+
38+
public KotlinSourceRemapper(StratumExt stratum) {
39+
this.stratum = stratum;
40+
}
41+
42+
@Override
43+
public int remapSourceLine(int line) {
44+
Pair<String, Integer> pair = stratum.getInputLine(line);
45+
if (pair == null || pair.getRight() == null) {
46+
return line;
47+
}
48+
return pair.getRight();
49+
}
50+
}
51+
}

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/symbol/SymbolExtractor.java‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import static com.datadog.debugger.instrumentation.ASMHelper.sortLocalVariables;
66

77
import com.datadog.debugger.instrumentation.ASMHelper;
8+
import datadog.trace.agent.tooling.stratum.SourceMap;
9+
import datadog.trace.agent.tooling.stratum.parser.Parser;
810
import datadog.trace.util.Strings;
911
import java.util.ArrayList;
1012
import java.util.Collection;
@@ -28,9 +30,11 @@
2830
import org.objectweb.asm.tree.LineNumberNode;
2931
import org.objectweb.asm.tree.LocalVariableNode;
3032
import org.objectweb.asm.tree.MethodNode;
33+
import org.slf4j.Logger;
3134
import org.slf4j.LoggerFactory;
3235

3336
public class SymbolExtractor {
37+
private static final Logger LOGGER = LoggerFactory.getLogger(SymbolExtractor.class);
3438

3539
public static Scope extract(byte[] classFileBuffer, String jarName) {
3640
ClassNode classNode = parseClassFile(classFileBuffer);
@@ -40,7 +44,16 @@ public static Scope extract(byte[] classFileBuffer, String jarName) {
4044
private static Scope extractScopes(ClassNode classNode, String jarName) {
4145
try {
4246
String sourceFile = extractSourceFile(classNode);
43-
List<Scope> methodScopes = extractMethods(classNode, sourceFile);
47+
SourceRemapper sourceRemapper = SourceRemapper.NOOP_REMAPPER;
48+
if (classNode.sourceDebug != null) {
49+
List<SourceMap> sourceMaps = Parser.parse(classNode.sourceDebug);
50+
if (sourceMaps.isEmpty()) {
51+
throw new IllegalStateException("No source maps found for " + classNode.name);
52+
}
53+
SourceMap sourceMap = sourceMaps.get(0);
54+
sourceRemapper = SourceRemapper.getSourceRemapper(classNode.sourceFile, sourceMap);
55+
}
56+
List<Scope> methodScopes = extractMethods(classNode, sourceFile, sourceRemapper);
4457
int classStartLine = Integer.MAX_VALUE;
4558
int classEndLine = 0;
4659
for (Scope scope : methodScopes) {
@@ -67,9 +80,8 @@ private static Scope extractScopes(ClassNode classNode, String jarName) {
6780
.scopes(new ArrayList<>(Collections.singletonList(classScope)))
6881
.build();
6982
} catch (Exception ex) {
70-
LoggerFactory.getLogger(SymbolExtractor.class)
71-
.debug(
72-
"Extracting scopes for class[{}] in jar[{}] failed: ", classNode.name, jarName, ex);
83+
LOGGER.debug(
84+
"Extracting scopes for class[{}] in jar[{}] failed: ", classNode.name, jarName, ex);
7385
return null;
7486
}
7587
}
@@ -102,10 +114,11 @@ private static List<Symbol> extractFields(ClassNode classNode) {
102114
return fields;
103115
}
104116

105-
private static List<Scope> extractMethods(ClassNode classNode, String sourceFile) {
117+
private static List<Scope> extractMethods(
118+
ClassNode classNode, String sourceFile, SourceRemapper sourceRemapper) {
106119
List<Scope> methodScopes = new ArrayList<>();
107120
for (MethodNode method : classNode.methods) {
108-
MethodLineInfo methodLineInfo = extractMethodLineInfo(method);
121+
MethodLineInfo methodLineInfo = extractMethodLineInfo(method, sourceRemapper);
109122
List<Scope> varScopes = new ArrayList<>();
110123
List<Symbol> methodSymbols = new ArrayList<>();
111124
int localVarBaseSlot = extractArgs(method, methodSymbols, methodLineInfo.start);
@@ -464,7 +477,8 @@ static List<Scope.LineRange> buildRanges(List<Integer> sortedLineNo) {
464477
return ranges;
465478
}
466479

467-
private static MethodLineInfo extractMethodLineInfo(MethodNode methodNode) {
480+
private static MethodLineInfo extractMethodLineInfo(
481+
MethodNode methodNode, SourceRemapper sourceRemapper) {
468482
Map<Label, Integer> map = new HashMap<>();
469483
List<Integer> lineNo = new ArrayList<>();
470484
Set<Integer> dedupSet = new HashSet<>();
@@ -473,10 +487,11 @@ private static MethodLineInfo extractMethodLineInfo(MethodNode methodNode) {
473487
while (node != null) {
474488
if (node.getType() == AbstractInsnNode.LINE) {
475489
LineNumberNode lineNumberNode = (LineNumberNode) node;
476-
if (dedupSet.add(lineNumberNode.line)) {
477-
lineNo.add(lineNumberNode.line);
490+
int newLine = sourceRemapper.remapSourceLine(lineNumberNode.line);
491+
if (dedupSet.add(newLine)) {
492+
lineNo.add(newLine);
478493
}
479-
maxLine = Math.max(lineNumberNode.line, maxLine);
494+
maxLine = Math.max(newLine, maxLine);
480495
}
481496
if (node.getType() == AbstractInsnNode.LABEL) {
482497
if (node instanceof LabelNode) {

‎dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/JvmLanguage.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package com.datadog.debugger.util;
2+
3+
public enum JvmLanguage {
4+
JAVA,
5+
KOTLIN,
6+
SCALA,
7+
GROOVY,
8+
UNKNOWN;
9+
10+
public static JvmLanguage of(String sourceFile) {
11+
if (sourceFile == null) {
12+
return UNKNOWN;
13+
}
14+
if (sourceFile.endsWith(".java")) {
15+
return JAVA;
16+
}
17+
if (sourceFile.endsWith(".kt")) {
18+
return KOTLIN;
19+
}
20+
if (sourceFile.endsWith(".scala")) {
21+
return SCALA;
22+
}
23+
if (sourceFile.endsWith(".groovy")) {
24+
return GROOVY;
25+
}
26+
return UNKNOWN;
27+
}
28+
}

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

Lines changed: 0 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -25,30 +25,18 @@
2525
import datadog.trace.bootstrap.debugger.ProbeId;
2626
import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
2727
import datadog.trace.bootstrap.debugger.util.Redaction;
28-
import java.io.File;
2928
import java.io.IOException;
3029
import java.lang.instrument.ClassFileTransformer;
3130
import java.lang.instrument.Instrumentation;
3231
import java.net.URISyntaxException;
33-
import java.net.URL;
34-
import java.net.URLClassLoader;
35-
import java.nio.file.Files;
36-
import java.nio.file.Path;
3732
import java.util.ArrayList;
3833
import java.util.Collection;
3934
import java.util.Collections;
40-
import java.util.Comparator;
4135
import java.util.HashMap;
4236
import java.util.Iterator;
4337
import java.util.List;
4438
import java.util.Map;
4539
import net.bytebuddy.agent.ByteBuddyAgent;
46-
import org.jetbrains.kotlin.cli.common.ExitCode;
47-
import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments;
48-
import org.jetbrains.kotlin.cli.common.messages.MessageRenderer;
49-
import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector;
50-
import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler;
51-
import org.jetbrains.kotlin.config.Services;
5240
import org.junit.jupiter.api.AfterEach;
5341
import org.junit.jupiter.api.Assertions;
5442
import org.junit.jupiter.api.BeforeEach;
@@ -415,49 +403,4 @@ public void instrumentationResult(ProbeDefinition definition, InstrumentationRes
415403
results.put(definition.getId(), result);
416404
}
417405
}
418-
419-
static class KotlinHelper {
420-
public static Class<?> compileAndLoad(
421-
String className, String sourceFileName, List<File> outputFilesToDelete) {
422-
K2JVMCompiler compiler = new K2JVMCompiler();
423-
K2JVMCompilerArguments args = compiler.createArguments();
424-
args.setFreeArgs(Collections.singletonList(sourceFileName));
425-
String compilerOutputDir = "/tmp/" + CapturedSnapshotTest.class.getSimpleName() + "-kotlin";
426-
args.setDestination(compilerOutputDir);
427-
args.setClasspath(System.getProperty("java.class.path"));
428-
ExitCode exitCode =
429-
compiler.exec(
430-
new PrintingMessageCollector(System.out, MessageRenderer.WITHOUT_PATHS, true),
431-
Services.EMPTY,
432-
args);
433-
434-
if (exitCode.getCode() != 0) {
435-
throw new RuntimeException("Kotlin compilation failed");
436-
}
437-
File compileOutputDirFile = new File(compilerOutputDir);
438-
try {
439-
URLClassLoader urlClassLoader =
440-
new URLClassLoader(new URL[] {compileOutputDirFile.toURI().toURL()});
441-
return urlClassLoader.loadClass(className);
442-
} catch (Exception ex) {
443-
throw new RuntimeException(ex);
444-
} finally {
445-
registerFilesToDeleteDir(compileOutputDirFile, outputFilesToDelete);
446-
}
447-
}
448-
449-
public static void registerFilesToDeleteDir(File dir, List<File> outputFilesToDelete) {
450-
if (!dir.exists()) {
451-
return;
452-
}
453-
try {
454-
Files.walk(dir.toPath())
455-
.sorted(Comparator.reverseOrder())
456-
.map(Path::toFile)
457-
.forEach(outputFilesToDelete::add);
458-
} catch (IOException ex) {
459-
ex.printStackTrace();
460-
}
461-
}
462-
}
463406
}

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package com.datadog.debugger.agent;
2+
3+
import java.io.File;
4+
import java.io.IOException;
5+
import java.net.URL;
6+
import java.net.URLClassLoader;
7+
import java.nio.file.Files;
8+
import java.nio.file.Path;
9+
import java.util.Collections;
10+
import java.util.Comparator;
11+
import java.util.List;
12+
import org.jetbrains.kotlin.cli.common.ExitCode;
13+
import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments;
14+
import org.jetbrains.kotlin.cli.common.messages.MessageRenderer;
15+
import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector;
16+
import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler;
17+
import org.jetbrains.kotlin.config.Services;
18+
19+
public class KotlinHelper {
20+
public static Class<?> compileAndLoad(
21+
String className, String sourceFileName, List<File> outputFilesToDelete) {
22+
K2JVMCompiler compiler = new K2JVMCompiler();
23+
K2JVMCompilerArguments args = compiler.createArguments();
24+
args.setFreeArgs(Collections.singletonList(sourceFileName));
25+
String compilerOutputDir = "/tmp/" + CapturedSnapshotTest.class.getSimpleName() + "-kotlin";
26+
args.setDestination(compilerOutputDir);
27+
args.setClasspath(System.getProperty("java.class.path"));
28+
ExitCode exitCode =
29+
compiler.exec(
30+
new PrintingMessageCollector(System.out, MessageRenderer.WITHOUT_PATHS, true),
31+
Services.EMPTY,
32+
args);
33+
34+
if (exitCode.getCode() != 0) {
35+
throw new RuntimeException("Kotlin compilation failed");
36+
}
37+
File compileOutputDirFile = new File(compilerOutputDir);
38+
try {
39+
URLClassLoader urlClassLoader =
40+
new URLClassLoader(new URL[] {compileOutputDirFile.toURI().toURL()});
41+
return urlClassLoader.loadClass(className);
42+
} catch (Exception ex) {
43+
throw new RuntimeException(ex);
44+
} finally {
45+
registerFilesToDeleteDir(compileOutputDirFile, outputFilesToDelete);
46+
}
47+
}
48+
49+
public static void registerFilesToDeleteDir(File dir, List<File> outputFilesToDelete) {
50+
if (!dir.exists()) {
51+
return;
52+
}
53+
try {
54+
Files.walk(dir.toPath())
55+
.sorted(Comparator.reverseOrder())
56+
.map(Path::toFile)
57+
.forEach(outputFilesToDelete::add);
58+
} catch (IOException ex) {
59+
ex.printStackTrace();
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)