Skip to content

Commit 5d109aa

Browse files
jpbempeljpbempel
authored
Optimize the SourceFile tracking (#8520)
We provide a new way to extract the SourceFile attribute from classfile that just scan the byte array to reconstruct the offsets of the constant pool and quickly fetch at the right location the content of the SourceFile attribute
1 parent a464b9b commit 5d109aa

20 files changed

+290
-78
lines changed

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

+131-1
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,20 @@
11
package com.datadog.debugger.util;
22

3+
import java.nio.charset.StandardCharsets;
34
import org.objectweb.asm.ClassReader;
45
import org.objectweb.asm.tree.ClassNode;
56

67
/** Helper class for extracting information of a class file */
78
public class ClassFileHelper {
9+
private static final int CONSTANT_POOL_COUNT_OFFSET = 8;
10+
private static final int CONSTANT_POOL_BASE_OFFSET = 10;
11+
812
public static String extractSourceFile(byte[] classFileBuffer) {
9-
// TODO maybe by scanning the byte array directly we can avoid doing an expensive parsing
13+
return extractSourceFileOffsetVersion(classFileBuffer);
14+
}
15+
16+
// Version using ASM library
17+
private static String extractSourceFileASM(byte[] classFileBuffer) {
1018
ClassReader classReader = new ClassReader(classFileBuffer);
1119
ClassNode classNode = new ClassNode();
1220
classReader.accept(classNode, ClassReader.SKIP_FRAMES);
@@ -33,4 +41,126 @@ public static String stripPackagePath(String classPath) {
3341
}
3442
return classPath;
3543
}
44+
45+
// Based on JVM spec: https://docs.oracle.com/javase/specs/jvms/se23/html/jvms-4.html
46+
// Extracts the SourceFile attribute from a Java class file byte array with minimal parsing.
47+
// This method is based on the JVM spec and does not use any external libraries.
48+
// We are scanning the constant pool to keep file offsets for later fetching of the SourceFile
49+
// attribute value. As the constant pool is a variable length structure, we need to scan them
50+
// and based on the tag, we can calculate the length of the entry to skip to the next one.
51+
private static String extractSourceFileOffsetVersion(byte[] classFileBytes) {
52+
// Quick validation of minimum class file size and magic number
53+
if (classFileBytes == null
54+
|| classFileBytes.length < 10
55+
|| classFileBytes[0] != (byte) 0xCA
56+
|| classFileBytes[1] != (byte) 0xFE
57+
|| classFileBytes[2] != (byte) 0xBA
58+
|| classFileBytes[3] != (byte) 0xBE) {
59+
return null;
60+
}
61+
int constantPoolCount = readUnsignedShort(classFileBytes, CONSTANT_POOL_COUNT_OFFSET);
62+
int[] constantPoolOffsets = new int[constantPoolCount];
63+
int currentOffset = CONSTANT_POOL_BASE_OFFSET;
64+
// based on the JVM spec, constant pool starts from index 1 until constantPoolCount - 1
65+
for (int i = 0; i < constantPoolCount - 1; i++) {
66+
constantPoolOffsets[i] = currentOffset;
67+
int tag = classFileBytes[constantPoolOffsets[i]];
68+
switch (tag) {
69+
case 1: // CONSTANT_Utf8
70+
int length = readUnsignedShort(classFileBytes, constantPoolOffsets[i] + 1);
71+
currentOffset += 3 + length;
72+
break;
73+
case 7: // CONSTANT_Class
74+
case 8: // CONSTANT_String
75+
case 16: // CONSTANT_MethodType
76+
case 19: // CONSTANT_Module
77+
case 20: // CONSTANT_Package
78+
currentOffset += 3;
79+
break;
80+
case 15: // CONSTANT_MethodHandle
81+
currentOffset += 4;
82+
break;
83+
case 3: // CONSTANT_Integer
84+
case 4: // CONSTANT_Float
85+
case 9: // CONSTANT_Fieldref
86+
case 10: // CONSTANT_Methodref
87+
case 11: // CONSTANT_InterfaceMethodref
88+
case 12: // CONSTANT_NameAndType
89+
case 17: // CONSTANT_Dynamic
90+
case 18: // CONSTANT_InvokeDynamic
91+
currentOffset += 5;
92+
break;
93+
case 5: // CONSTANT_Long
94+
case 6: // CONSTANT_Double
95+
currentOffset += 9;
96+
i++; // Double slot
97+
break;
98+
default:
99+
throw new IllegalArgumentException("Unknown constant pool tag: " + tag);
100+
}
101+
}
102+
currentOffset += 2; // Skip access flags
103+
currentOffset += 2; // Skip this class
104+
currentOffset += 2; // Skip super class
105+
int interfacesCount = readUnsignedShort(classFileBytes, currentOffset);
106+
currentOffset += 2 + interfacesCount * 2; // Skip interfaces
107+
// skip fields
108+
currentOffset = skipFieldsOrMethods(classFileBytes, currentOffset);
109+
// skip Methods
110+
currentOffset = skipFieldsOrMethods(classFileBytes, currentOffset);
111+
int attributesCount = readUnsignedShort(classFileBytes, currentOffset);
112+
currentOffset += 2; // Skip attributes count
113+
for (int i = 0; i < attributesCount; i++) {
114+
int attributeNameIndex = readUnsignedShort(classFileBytes, currentOffset);
115+
currentOffset += 2; // Skip attribute name index
116+
int attributeLength = (int) readUnsignedInt(classFileBytes, currentOffset);
117+
currentOffset += 4; // Skip attribute length
118+
if (attributeNameIndex == 0) {
119+
continue;
120+
}
121+
// read attribute name
122+
int utf8Offset = constantPoolOffsets[attributeNameIndex - 1];
123+
int utf8Len = readUnsignedShort(classFileBytes, utf8Offset + 1);
124+
String utf8 = new String(classFileBytes, utf8Offset + 3, utf8Len, StandardCharsets.UTF_8);
125+
if ("SourceFile".equals(utf8)) {
126+
// read SourceFile attribute
127+
int sourceFileIndex = readUnsignedShort(classFileBytes, currentOffset);
128+
int sourceFileOffset = constantPoolOffsets[sourceFileIndex - 1];
129+
int sourceFileLen = readUnsignedShort(classFileBytes, sourceFileOffset + 1);
130+
return new String(
131+
classFileBytes, sourceFileOffset + 3, sourceFileLen, StandardCharsets.UTF_8);
132+
}
133+
currentOffset += attributeLength; // Skip attribute data
134+
}
135+
return null;
136+
}
137+
138+
private static int skipFieldsOrMethods(byte[] classFileBytes, int currentOffset) {
139+
int fieldsCount = readUnsignedShort(classFileBytes, currentOffset);
140+
currentOffset += 2; // Skip count
141+
for (int i = 0; i < fieldsCount; i++) {
142+
currentOffset += 6; // Skip access flags, name index, descriptor index
143+
int attributesCount = readUnsignedShort(classFileBytes, currentOffset);
144+
currentOffset += 2; // Skip attributes count
145+
for (int j = 0; j < attributesCount; j++) {
146+
currentOffset += 2; // Skip attribute name index
147+
int attributeLength = (int) readUnsignedInt(classFileBytes, currentOffset);
148+
currentOffset += 4 + attributeLength; // Skip attribute length and data
149+
}
150+
}
151+
return currentOffset;
152+
}
153+
154+
// read unsigned short from byte array
155+
private static int readUnsignedShort(byte[] bytes, int offset) {
156+
return ((bytes[offset] & 0xFF) << 8) | (bytes[offset + 1] & 0xFF);
157+
}
158+
159+
// read unsigned int from byte array
160+
private static long readUnsignedInt(byte[] bytes, int offset) {
161+
return ((long) (bytes[offset] & 0xFF) << 24)
162+
+ ((bytes[offset + 1] & 0xFF) << 16)
163+
+ ((bytes[offset + 2] & 0xFF) << 8)
164+
+ (bytes[offset + 3] & 0xFF);
165+
}
36166
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import static com.datadog.debugger.util.MoshiSnapshotHelper.REDACTED_IDENT_REASON;
88
import static com.datadog.debugger.util.MoshiSnapshotHelper.REDACTED_TYPE_REASON;
99
import static com.datadog.debugger.util.MoshiSnapshotTestHelper.VALUE_ADAPTER;
10-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
1110
import static datadog.trace.bootstrap.debugger.util.Redaction.REDACTED_VALUE;
1211
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
1312
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -27,6 +26,7 @@
2726
import static utils.InstrumentationTestHelper.getLineForLineProbe;
2827
import static utils.InstrumentationTestHelper.loadClass;
2928
import static utils.TestHelper.getFixtureContent;
29+
import static utils.TestHelper.setFieldInConfig;
3030

3131
import com.datadog.debugger.el.DSL;
3232
import com.datadog.debugger.el.ProbeCondition;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33
import static com.datadog.debugger.util.MoshiSnapshotHelper.NOT_CAPTURED_REASON;
44
import static com.datadog.debugger.util.MoshiSnapshotTestHelper.VALUE_ADAPTER;
5-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
65
import static org.junit.jupiter.api.Assertions.assertEquals;
76
import static org.junit.jupiter.api.Assertions.assertNotNull;
87
import static org.junit.jupiter.api.Assertions.assertTrue;
98
import static org.mockito.Mockito.mock;
9+
import static utils.TestHelper.setFieldInConfig;
1010

1111
import com.datadog.debugger.instrumentation.InstrumentationResult;
1212
import com.datadog.debugger.probe.LogProbe;

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package com.datadog.debugger.agent;
22

3-
import static com.datadog.debugger.util.TestHelper.setEnvVar;
4-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
53
import static org.junit.jupiter.api.Assertions.assertEquals;
64
import static org.junit.jupiter.api.Assertions.assertNotNull;
75
import static org.junit.jupiter.api.Assumptions.assumeTrue;
@@ -14,6 +12,8 @@
1412
import static org.mockito.Mockito.never;
1513
import static org.mockito.Mockito.verify;
1614
import static org.mockito.Mockito.when;
15+
import static utils.TestHelper.setEnvVar;
16+
import static utils.TestHelper.setFieldInConfig;
1717

1818
import com.datadog.debugger.util.RemoteConfigHelper;
1919
import datadog.common.container.ContainerInfo;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package com.datadog.debugger.agent;
22

3-
import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
43
import static org.junit.jupiter.api.Assertions.assertEquals;
54
import static org.junit.jupiter.api.Assertions.assertNull;
65
import static org.mockito.ArgumentMatchers.any;
@@ -10,6 +9,7 @@
109
import static org.mockito.Mockito.times;
1110
import static org.mockito.Mockito.verify;
1211
import static org.mockito.Mockito.when;
12+
import static utils.TestClassFileHelper.getClassFileBytes;
1313

1414
import com.datadog.debugger.instrumentation.DiagnosticMessage;
1515
import com.datadog.debugger.instrumentation.InstrumentationResult;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package com.datadog.debugger.agent;
22

3-
import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
43
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static utils.TestClassFileHelper.getClassFileBytes;
55

66
import com.datadog.debugger.probe.LogProbe;
77
import java.lang.instrument.IllegalClassFormatException;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
package com.datadog.debugger.agent;
22

3-
import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
43
import static java.util.Arrays.asList;
54
import static org.junit.jupiter.api.Assertions.*;
5+
import static utils.TestClassFileHelper.getClassFileBytes;
66

77
import com.datadog.debugger.probe.LogProbe;
88
import com.datadog.debugger.probe.MetricProbe;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package com.datadog.debugger.exception;
22

33
import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
4-
import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
54
import static java.util.Collections.emptyList;
65
import static java.util.Collections.singletonList;
76
import static java.util.stream.Collectors.toList;
@@ -16,6 +15,7 @@
1615
import static org.mockito.Mockito.times;
1716
import static org.mockito.Mockito.verify;
1817
import static org.mockito.Mockito.when;
18+
import static utils.TestHelper.assertWithTimeout;
1919

2020
import com.datadog.debugger.agent.ConfigurationAcceptor;
2121
import com.datadog.debugger.agent.ConfigurationUpdater;

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,13 @@
66
import static com.datadog.debugger.exception.DefaultExceptionDebugger.ERROR_DEBUG_INFO_CAPTURED;
77
import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
88
import static com.datadog.debugger.util.MoshiSnapshotTestHelper.getValue;
9-
import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
10-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
119
import static org.junit.jupiter.api.Assertions.assertEquals;
1210
import static org.junit.jupiter.api.Assertions.assertTrue;
1311
import static org.mockito.Mockito.mock;
1412
import static org.mockito.Mockito.when;
1513
import static utils.InstrumentationTestHelper.compileAndLoadClass;
14+
import static utils.TestHelper.assertWithTimeout;
15+
import static utils.TestHelper.setFieldInConfig;
1616

1717
import com.datadog.debugger.agent.ClassesToRetransformFinder;
1818
import com.datadog.debugger.agent.Configuration;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package com.datadog.debugger.origin;
22

3-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
43
import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_FRAME;
54
import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_PREFIX;
65
import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_TYPE;
@@ -14,6 +13,7 @@
1413
import static org.junit.jupiter.api.Assertions.assertNull;
1514
import static org.junit.jupiter.api.Assertions.assertTrue;
1615
import static utils.InstrumentationTestHelper.compileAndLoadClass;
16+
import static utils.TestHelper.setFieldInConfig;
1717

1818
import com.datadog.debugger.agent.CapturingTestBase;
1919
import com.datadog.debugger.codeorigin.DefaultCodeOriginRecorder;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
package com.datadog.debugger.trigger;
22

33
import static com.datadog.debugger.el.DSL.*;
4-
import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
54
import static java.lang.String.format;
65
import static org.junit.jupiter.api.Assertions.assertEquals;
76
import static org.junit.jupiter.api.Assertions.assertTrue;
87
import static org.mockito.ArgumentMatchers.eq;
98
import static org.mockito.Mockito.verify;
109
import static utils.InstrumentationTestHelper.compileAndLoadClass;
10+
import static utils.TestHelper.setFieldInConfig;
1111

1212
import com.datadog.debugger.agent.CapturingTestBase;
1313
import com.datadog.debugger.agent.Configuration;

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

+28-16
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,37 @@
11
package com.datadog.debugger.util;
22

3-
import static datadog.trace.util.Strings.getResourceName;
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
44

5-
import java.io.ByteArrayOutputStream;
65
import java.io.IOException;
7-
import java.io.InputStream;
8-
import java.net.URL;
6+
import java.net.URISyntaxException;
7+
import java.nio.file.Files;
8+
import java.nio.file.Paths;
9+
import org.junit.jupiter.api.Test;
910

1011
public class ClassFileHelperTest {
11-
public static byte[] getClassFileBytes(Class<?> clazz) {
12-
URL resource = clazz.getResource("/" + getResourceName(clazz.getTypeName()));
13-
byte[] buffer = new byte[4096];
14-
ByteArrayOutputStream os = new ByteArrayOutputStream();
15-
try (InputStream is = resource.openStream()) {
16-
int readBytes;
17-
while ((readBytes = is.read(buffer)) != -1) {
18-
os.write(buffer, 0, readBytes);
19-
}
20-
} catch (IOException e) {
21-
e.printStackTrace();
12+
13+
@Test
14+
public void extractSourceFile() {
15+
assertEquals(
16+
"JDK8.java",
17+
ClassFileHelper.extractSourceFile(
18+
readClassFileBytes("/com/datadog/debugger/classfiles/JDK8.class")));
19+
assertEquals(
20+
"JDK23.java",
21+
ClassFileHelper.extractSourceFile(
22+
readClassFileBytes("/com/datadog/debugger/classfiles/JDK23.class")));
23+
// big classfile (80KB)
24+
assertEquals(
25+
"CommandLine.java",
26+
ClassFileHelper.extractSourceFile(
27+
readClassFileBytes("/com/datadog/debugger/classfiles/CommandLine.class")));
28+
}
29+
30+
private static byte[] readClassFileBytes(String fileName) {
31+
try {
32+
return Files.readAllBytes(Paths.get(ClassFileHelperTest.class.getResource(fileName).toURI()));
33+
} catch (IOException | URISyntaxException e) {
34+
throw new RuntimeException(e);
2235
}
23-
return os.toByteArray();
2436
}
2537
}

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

-49
This file was deleted.

0 commit comments

Comments
 (0)