From f2e4b7d4b11628e42e023e60887fabba04e9e589 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Wed, 12 Dec 2018 20:42:56 -0800 Subject: [PATCH] Use javaagent to determine if logger class is initialized in test --- .../trace/agent/CustomLogManagerTest.groovy | 3 +- .../agent/test/IntegrationTestUtils.java | 15 ++++++- .../jvmbootstraptest/LogManagerSetter.java | 40 ++++++++++++++----- 3 files changed, 46 insertions(+), 12 deletions(-) diff --git a/dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/CustomLogManagerTest.groovy b/dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/CustomLogManagerTest.groovy index ec5ec4b464..d8cfa0704e 100644 --- a/dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/CustomLogManagerTest.groovy +++ b/dd-java-agent-ittests/src/test/groovy/datadog/trace/agent/CustomLogManagerTest.groovy @@ -21,11 +21,12 @@ class CustomLogManagerTest extends Specification { break } } + final URL customAgent = IntegrationTestUtils.createJarWithClasses(LogManagerSetter.getName(), LogManagerSetter) expect: agentArg != null IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName() - , [agentArg, "-Ddd.jmxfetch.enabled=true"] as String[] + , [agentArg, "-javaagent:" + customAgent.getPath(), "-Ddd.jmxfetch.enabled=true"] as String[] , "" as String[] , true) == 0 } diff --git a/dd-java-agent-ittests/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java b/dd-java-agent-ittests/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java index a89760f606..737e5e9304 100644 --- a/dd-java-agent-ittests/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java +++ b/dd-java-agent-ittests/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java @@ -19,6 +19,7 @@ import java.util.Arrays; import java.util.List; import java.util.UUID; import java.util.concurrent.Callable; +import java.util.jar.Attributes; import java.util.jar.JarEntry; import java.util.jar.JarOutputStream; import java.util.jar.Manifest; @@ -77,20 +78,32 @@ public class IntegrationTestUtils { } } + /** See {@link IntegrationTestUtils#createJarWithClasses(String, Class[])} */ + public static URL createJarWithClasses(final Class... classes) throws IOException { + return createJarWithClasses(null, classes); + } /** * Create a temporary jar on the filesystem with the bytes of the given classes. * *

The jar file will be removed when the jvm exits. * + * @param mainClassname The name of the class to use for Main-Class and Premain-Class. May be null * @param classes classes to package into the jar. * @return the location of the newly created jar. * @throws IOException */ - public static URL createJarWithClasses(final Class... classes) throws IOException { + public static URL createJarWithClasses(final String mainClassname, final Class... classes) + throws IOException { final File tmpJar = File.createTempFile(UUID.randomUUID().toString() + "-", ".jar"); tmpJar.deleteOnExit(); final Manifest manifest = new Manifest(); + if (mainClassname != null) { + Attributes mainAttributes = manifest.getMainAttributes(); + mainAttributes.put(Attributes.Name.MANIFEST_VERSION, "1.0"); + mainAttributes.put(Attributes.Name.MAIN_CLASS, mainClassname); + mainAttributes.put(new Attributes.Name("Premain-Class"), mainClassname); + } final JarOutputStream target = new JarOutputStream(new FileOutputStream(tmpJar), manifest); for (final Class clazz : classes) { addToJar(clazz, target); diff --git a/dd-java-agent-ittests/src/test/java/jvmbootstraptest/LogManagerSetter.java b/dd-java-agent-ittests/src/test/java/jvmbootstraptest/LogManagerSetter.java index d8f297980b..cd2dc3f949 100644 --- a/dd-java-agent-ittests/src/test/java/jvmbootstraptest/LogManagerSetter.java +++ b/dd-java-agent-ittests/src/test/java/jvmbootstraptest/LogManagerSetter.java @@ -1,6 +1,10 @@ package jvmbootstraptest; -import java.lang.reflect.Method; +import java.lang.instrument.ClassFileTransformer; +import java.lang.instrument.IllegalClassFormatException; +import java.lang.instrument.Instrumentation; +import java.security.ProtectionDomain; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.LogManager; public class LogManagerSetter { @@ -8,22 +12,38 @@ public class LogManagerSetter { // loggerClassName = java.util.logging.Logger private static final String loggerClassName = "java.util.logging.TMP".replaceFirst("TMP", "Logger"); + private static final String loggerInternalName = loggerClassName.replace('.', '/'); + private static final AtomicBoolean loggerInitialized = new AtomicBoolean(false); + + public static void premain(final String agentArgs, final Instrumentation inst) throws Exception { + inst.addTransformer( + new ClassFileTransformer() { + @Override + public byte[] transform( + ClassLoader loader, + String className, + Class classBeingRedefined, + ProtectionDomain protectionDomain, + byte[] classfileBuffer) + throws IllegalClassFormatException { + if (loggerInternalName.equals(className)) { + loggerInitialized.compareAndSet(false, true); + } + return null; + } + }); + } public static void main(String... args) throws Exception { final ClassLoader systemLoader = ClassLoader.getSystemClassLoader(); // once the logger class has been initialized, we know jmxfetch is running and can proceed with // the test - final Method method = - ClassLoader.class.getDeclaredMethod("findBootstrapClassOrNull", String.class); - try { - method.setAccessible(true); - while (method.invoke(systemLoader, loggerClassName) == null) { - Thread.sleep(1); - } - } finally { - method.setAccessible(false); + while (!loggerInitialized.get()) { + Thread.sleep(1); } + systemLoader.loadClass(loggerClassName); + // at this point the logger is loaded and fully initialized System.setProperty("java.util.logging.manager", CustomLogManager.class.getName()); customAssert(