Enhance Support for spring-boot classloader

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/ClassLoaderTest.java

@@ -0,0 +1,103 @@
+package com.datadoghq.agent;
+
+import com.datadoghq.trace.Trace;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.reflect.Method;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.UUID;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
+import java.util.jar.Manifest;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class ClassLoaderTest {
+
+  /** Assert that we can instrument classloaders which cannot resolve agent advice classes. */
+  @Test
+  public void instrumentClassLoadersWithoutAgentClasses() throws Exception {
+    URL[] classpath = new URL[] {createJarWithClasses(ClassToInstrument.class, Trace.class)};
+    URLClassLoader loader = new URLClassLoader(classpath, null);
+
+    try {
+      loader.loadClass("com.datadoghq.agent.TracingAgent");
+      Assert.fail("loader should not see agent classes.");
+    } catch (ClassNotFoundException cnfe) {
+      // Good. loader can't see agent classes.
+    }
+
+    Class<?> instrumentedClass = loader.loadClass(ClassToInstrument.class.getName());
+    Assert.assertEquals(
+        "Class must be loaded by loader.", loader, instrumentedClass.getClassLoader());
+
+    final Class<?> rulesManagerClass =
+        Class.forName(
+            "com.datadoghq.agent.InstrumentationRulesManager",
+            true,
+            ClassLoader.getSystemClassLoader());
+    Method isRegisteredMethod = rulesManagerClass.getMethod("isRegistered", Object.class);
+    Assert.assertTrue(
+        "Agent did not initialized loader.", (boolean) isRegisteredMethod.invoke(null, loader));
+    loader.close();
+  }
+
+  /** com.foo.Bar -> com/foo/Bar.class */
+  public static String getResourceName(Class<?> clazz) {
+    return clazz.getName().replace('.', '/') + ".class";
+  }
+
+  /**
+   * Create a temporary jar on the filesystem with the bytes of the given classes.
+   *
+   * <p>The jar file will be removed when the jvm exits.
+   *
+   * @param classes classes to package into the jar.
+   * @return the location of the newly created jar.
+   * @throws IOException
+   */
+  public static URL createJarWithClasses(Class<?>... classes) throws IOException {
+    final File tmpJar = File.createTempFile(UUID.randomUUID() + "", ".jar");
+    tmpJar.deleteOnExit();
+
+    final Manifest manifest = new Manifest();
+    JarOutputStream target = new JarOutputStream(new FileOutputStream(tmpJar), manifest);
+    for (Class<?> clazz : classes) {
+      addToJar(clazz, target);
+    }
+    target.close();
+
+    return tmpJar.toURI().toURL();
+  }
+
+  private static void addToJar(Class<?> clazz, JarOutputStream jarOutputStream) throws IOException {
+    InputStream inputStream = null;
+    try {
+      JarEntry entry = new JarEntry(getResourceName(clazz));
+      jarOutputStream.putNextEntry(entry);
+      inputStream = clazz.getClassLoader().getResourceAsStream(getResourceName(clazz));
+
+      byte[] buffer = new byte[1024];
+      while (true) {
+        int count = inputStream.read(buffer);
+        if (count == -1) {
+          break;
+        }
+        jarOutputStream.write(buffer, 0, count);
+      }
+      jarOutputStream.closeEntry();
+    } finally {
+      if (inputStream != null) {
+        inputStream.close();
+      }
+    }
+  }
+
+  public static class ClassToInstrument {
+    @Trace
+    public static void someMethod() {}
+  }
+}

dd-java-agent/integrations/apache-httpclient-4.3/src/main/java/dd/inst/apachehttpclient/ApacheHttpClientInstrumentation.java

@@ -1,11 +1,11 @@
 package dd.inst.apachehttpclient;
 
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.*;
 
 import com.datadoghq.agent.integration.DDTracingClientExec;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.util.GlobalTracer;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -32,11 +32,10 @@ public AgentBuilder instrument(AgentBuilder agentBuilder) {
                 "org.apache.http.conn.routing.HttpRoute",
                 "org.apache.http.impl.execchain.ClientExecChain"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     isMethod().and(named("decorateProtocolExec")),
-                    ApacheHttpClientAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    ApacheHttpClientAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/aws-sdk/src/main/java/dd/inst/aws/AWSClientInstrumentation.java

@@ -1,14 +1,14 @@
 package dd.inst.aws;
 
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.amazonaws.client.builder.AwsClientBuilder;
 import com.amazonaws.handlers.RequestHandler2;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.contrib.aws.TracingRequestHandler;
 import io.opentracing.util.GlobalTracer;
@@ -30,11 +30,10 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
                 "com.amazonaws.http.client.HttpClientFactory",
                 "com.amazonaws.http.apache.utils.ApacheUtils"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("build").and(takesArguments(0)).and(isPublic()),
-                    AWSClientAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    AWSClientAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/datastax-cassandra-3.2/src/main/java/dd/inst/datastax/cassandra/CassandraClientInstrumentation.java

@@ -1,11 +1,11 @@
 package dd.inst.datastax.cassandra;
 
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.*;
 
 import com.datastax.driver.core.Session;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.Tracer;
 import io.opentracing.util.GlobalTracer;
@@ -36,11 +36,10 @@ public AgentBuilder instrument(AgentBuilder agentBuilder) {
                 "com.google.common.util.concurrent.Futures",
                 "com.google.common.util.concurrent.ListenableFuture"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     isMethod().and(isPrivate()).and(named("newSession")).and(takesArguments(0)),
-                    CassandraClientAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    CassandraClientAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-1/src/main/java/dd/inst/jms1/JMS1MessageConsumerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.SpanContext;
@@ -36,14 +36,13 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageConsumer"))),
             not(classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener")))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("receive").and(takesArguments(0)).and(isPublic()),
                     ConsumerAdvice.class.getName())
                 .advice(
                     named("receiveNoWait").and(takesArguments(0)).and(isPublic()),
-                    ConsumerAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    ConsumerAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-1/src/main/java/dd/inst/jms1/JMS1MessageListenerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.SpanContext;
@@ -35,13 +35,12 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageListener"))),
             not(classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener")))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("onMessage")
                         .and(takesArgument(0, named("javax.jms.Message")))
                         .and(isPublic()),
-                    MessageListenerAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    MessageListenerAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-1/src/main/java/dd/inst/jms1/JMS1MessageProducerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.propagation.Format;
@@ -36,7 +36,7 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageProducer"))),
             not(classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener")))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("send").and(takesArgument(0, named("javax.jms.Message"))).and(isPublic()),
                     ProducerAdvice.class.getName())
@@ -45,8 +45,7 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
                         .and(takesArgument(0, named("javax.jms.Destination")))
                         .and(takesArgument(1, named("javax.jms.Message")))
                         .and(isPublic()),
-                    ProducerWithDestinationAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    ProducerWithDestinationAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-2/src/main/java/dd/inst/jms2/JMS2MessageConsumerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.SpanContext;
@@ -36,14 +36,13 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageConsumer"))),
             classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("receive").and(takesArguments(0)).and(isPublic()),
                     ConsumerAdvice.class.getName())
                 .advice(
                     named("receiveNoWait").and(takesArguments(0)).and(isPublic()),
-                    ConsumerAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    ConsumerAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-2/src/main/java/dd/inst/jms2/JMS2MessageListenerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.SpanContext;
@@ -35,13 +35,12 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageListener"))),
             classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("onMessage")
                         .and(takesArgument(0, named("javax.jms.Message")))
                         .and(isPublic()),
-                    MessageListenerAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    MessageListenerAdvice.class.getName()))
         .asDecorator();
   }
 

dd-java-agent/integrations/jms-2/src/main/java/dd/inst/jms2/JMS2MessageProducerInstrumentation.java

@@ -2,7 +2,6 @@
 
 import static com.datadoghq.agent.integration.JmsUtil.toResourceName;
 import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
-import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -13,6 +12,7 @@
 import com.datadoghq.agent.integration.MessagePropertyTextMap;
 import com.datadoghq.trace.DDTags;
 import com.google.auto.service.AutoService;
+import dd.trace.DDAdvice;
 import dd.trace.Instrumenter;
 import io.opentracing.ActiveSpan;
 import io.opentracing.propagation.Format;
@@ -36,7 +36,7 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
             not(isInterface()).and(hasSuperType(named("javax.jms.MessageProducer"))),
             classLoaderHasClasses("javax.jms.JMSContext", "javax.jms.CompletionListener"))
         .transform(
-            new AgentBuilder.Transformer.ForAdvice()
+            DDAdvice.create()
                 .advice(
                     named("send").and(takesArgument(0, named("javax.jms.Message"))).and(isPublic()),
                     ProducerAdvice.class.getName())
@@ -45,8 +45,7 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
                         .and(takesArgument(0, named("javax.jms.Destination")))
                         .and(takesArgument(1, named("javax.jms.Message")))
                         .and(isPublic()),
-                    ProducerWithDestinationAdvice.class.getName())
-                .withExceptionHandler(defaultExceptionHandler()))
+                    ProducerWithDestinationAdvice.class.getName()))
         .asDecorator();
   }