Skip to content

Commit 1458290

Browse files
mccullsmcculls
authored
Load instrumentation classes in separate classloader so they can be unloaded after use (DataDog#4484)
1 parent 475c0d2 commit 1458290

5 files changed

Lines changed: 120 additions & 1 deletion

File tree

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

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import java.io.File;
44
import java.io.IOException;
55
import java.io.InputStream;
6+
import java.lang.ref.WeakReference;
67
import java.net.URL;
78
import java.security.CodeSource;
89
import java.security.SecureClassLoader;
@@ -31,6 +32,10 @@ public final class DatadogClassLoader extends SecureClassLoader {
3132
private final String agentResourcePrefix;
3233
private final AgentJarIndex agentJarIndex;
3334

35+
private final Object instrumentationClassLoaderLock = new Object();
36+
private volatile WeakReference<InstrumentationClassLoader> instrumentationClassLoader =
37+
new WeakReference<>(new InstrumentationClassLoader(this));
38+
3439
public DatadogClassLoader(final URL agentJarURL, final ClassLoader parent) throws Exception {
3540
super(parent);
3641

@@ -85,8 +90,33 @@ protected Enumeration<URL> findResources(String name) {
8590
return Collections.emptyEnumeration();
8691
}
8792

93+
@Override
94+
protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
95+
if (name.startsWith("datadog.trace.instrumentation.")
96+
&& (name.endsWith("$Muzzle") || name.endsWith("Instrumentation"))) {
97+
InstrumentationClassLoader cl;
98+
if (null == (cl = instrumentationClassLoader.get())) {
99+
synchronized (instrumentationClassLoaderLock) {
100+
if (null == (cl = instrumentationClassLoader.get())) {
101+
// previous instance was unloaded, create fresh one
102+
cl = new InstrumentationClassLoader(this);
103+
instrumentationClassLoader = new WeakReference<>(cl);
104+
}
105+
}
106+
}
107+
return cl.loadInstrumentationClass(name, agentCodeSource);
108+
} else {
109+
return super.loadClass(name, resolve);
110+
}
111+
}
112+
88113
@Override
89114
protected Class<?> findClass(String name) throws ClassNotFoundException {
115+
byte[] buf = loadClassBytes(name);
116+
return defineClass(name, buf, 0, buf.length, agentCodeSource);
117+
}
118+
119+
byte[] loadClassBytes(String name) throws ClassNotFoundException {
90120
String entryName = agentJarIndex.classEntryName(name);
91121
if (null != entryName) {
92122
JarEntry jarEntry = agentJarFile.getJarEntry(entryName);
@@ -102,7 +132,7 @@ protected Class<?> findClass(String name) throws ClassNotFoundException {
102132
bytesRead += delta;
103133
}
104134
if (bytesRead == buf.length) {
105-
return defineClass(name, buf, 0, buf.length, agentCodeSource);
135+
return buf;
106136
} else {
107137
log.warn("Malformed class data at {}", jarEntry);
108138
}

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package datadog.trace.bootstrap;
2+
3+
import java.security.CodeSource;
4+
import java.security.SecureClassLoader;
5+
6+
/** Holds Muzzle and Instrumentation classes, so they can be unloaded separately to the agent. */
7+
final class InstrumentationClassLoader extends SecureClassLoader {
8+
static {
9+
ClassLoader.registerAsParallelCapable();
10+
}
11+
12+
InstrumentationClassLoader(DatadogClassLoader parent) {
13+
super(parent);
14+
}
15+
16+
Class<?> loadInstrumentationClass(String name, CodeSource agentCodeSource)
17+
throws ClassNotFoundException {
18+
synchronized (getClassLoadingLock(name)) {
19+
Class<?> instrumentationClass = findLoadedClass(name);
20+
if (null != instrumentationClass) {
21+
return instrumentationClass;
22+
} else {
23+
// load bytecode from dd-java-agent jar, but define the class locally
24+
byte[] buf = ((DatadogClassLoader) getParent()).loadClassBytes(name);
25+
return defineClass(name, buf, 0, buf.length, agentCodeSource);
26+
}
27+
}
28+
}
29+
30+
@Override
31+
public String toString() {
32+
return "instrumentation";
33+
}
34+
}

‎dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
6868
+ "datadog.trace.bootstrap.BootstrapProxy:build_time,"
6969
+ "datadog.trace.bootstrap.CallDepthThreadLocalMap:build_time,"
7070
+ "datadog.trace.bootstrap.DatadogClassLoader:build_time,"
71+
+ "datadog.trace.bootstrap.InstrumentationClassLoader:build_time,"
7172
+ "datadog.trace.bootstrap.FieldBackedContextStores:build_time,"
7273
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState:build_time,"
7374
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter:build_time,"

‎dd-java-agent/src/test/groovy/datadog/trace/agent/InstrumenterUnloadTest.groovy‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package datadog.trace.agent
2+
3+
import datadog.trace.agent.test.IntegrationTestUtils
4+
import jvmbootstraptest.UnloadingChecker
5+
import spock.lang.Specification
6+
import spock.lang.Timeout
7+
8+
@Timeout(30)
9+
class InstrumenterUnloadTest extends Specification {
10+
11+
private static final String DEFAULT_LOG_LEVEL = "debug"
12+
private static final String API_KEY = "01234567890abcdef123456789ABCDEF"
13+
14+
// Run test using forked jvm
15+
def "instrumenter and muzzle classes can be unloaded after use"() {
16+
setup:
17+
def testOutput = new ByteArrayOutputStream()
18+
19+
when:
20+
int returnCode = IntegrationTestUtils.runOnSeparateJvm(UnloadingChecker.getName()
21+
, [
22+
"-verbose:class",
23+
"-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL"
24+
] as String[]
25+
, "" as String[]
26+
, ["DD_API_KEY": API_KEY]
27+
, new PrintStream(testOutput))
28+
29+
int unloadCount = 0
30+
new ByteArrayInputStream((testOutput.toByteArray())).eachLine {
31+
System.out.println(it)
32+
if (it =~ /(?i)unload.* datadog.trace.instrumentation./) {
33+
unloadCount++
34+
}
35+
}
36+
37+
then:
38+
returnCode == 0
39+
unloadCount > 0
40+
}
41+
}

‎dd-java-agent/src/test/java/jvmbootstraptest/UnloadingChecker.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package jvmbootstraptest;
2+
3+
import datadog.trace.test.util.GCUtils;
4+
5+
public class UnloadingChecker {
6+
public static void main(final String[] args) {
7+
try {
8+
GCUtils.awaitGC();
9+
} catch (InterruptedException e) {
10+
e.printStackTrace();
11+
}
12+
}
13+
}

0 commit comments

Comments
 (0)