[rmi] Instrumentation for RMI 1.2 and later

includes instrumentation of client and server endpoints
currently missing passing of execution context from client to server

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -92,6 +92,7 @@ public class AgentInstaller {
                         not(
                             named("java.net.URL")
                                 .or(named("java.net.HttpURLConnection"))
+                                .or(nameStartsWith("java.rmi."))
                                 .or(nameStartsWith("java.util.concurrent."))
                                 .or(
                                     nameStartsWith("java.util.logging.")
@@ -111,6 +112,7 @@ public class AgentInstaller {
                     .and(
                         not(
                             nameStartsWith("sun.net.www.protocol.")
+                                .or(nameStartsWith("sun.rmi.server."))
                                 .or(named("sun.net.www.http.HttpClient")))))
             .or(nameStartsWith("jdk."))
             .or(nameStartsWith("org.aspectj."))

dd-java-agent/instrumentation/rmi/rmi.gradle

@@ -0,0 +1,12 @@
+muzzle {
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+task "rmic", dependsOn: testClasses {
+  def clazz = 'rmi.app.ServerLegacy'
+  String command = """rmic -g -keep -classpath ${sourceSets.test.output.classesDirs.asPath} -d ${buildDir}/classes/java/test ${clazz}"""
+  command.execute().text
+}
+
+test.dependsOn "rmic"

dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/ClientAdvice.java

@@ -0,0 +1,41 @@
+package datadog.trace.instrumentation.rmi;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+
+import datadog.trace.api.DDTags;
+import datadog.trace.instrumentation.api.AgentScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import java.lang.reflect.Method;
+import net.bytebuddy.asm.Advice;
+
+public class ClientAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(@Advice.Argument(value = 1) final Method method) {
+    if (activeSpan() == null) {
+      return null;
+    }
+    final AgentSpan span =
+        startSpan("rmi.invoke")
+            .setTag(
+                DDTags.RESOURCE_NAME,
+                method.getDeclaringClass().getSimpleName() + "#" + method.getName())
+            .setTag("span.origin.type", method.getDeclaringClass().getCanonicalName());
+    return activateSpan(span, true);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void stopSpan(
+      @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+    if (scope == null) {
+      return;
+    }
+    final AgentSpan span = scope.span();
+    if (throwable != null) {
+      span.setError(true);
+      span.addThrowable(throwable);
+    }
+    scope.close();
+  }
+}

dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/RmiClientInstrumentation.java

@@ -0,0 +1,39 @@
+package datadog.trace.instrumentation.rmi;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class RmiClientInstrumentation extends Instrumenter.Default {
+
+  public RmiClientInstrumentation() {
+    super("rmi", "rmi-client");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("sun.rmi.server.UnicastRef")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("invoke"))
+            .and(takesArgument(0, named("java.rmi.Remote")))
+            .and(takesArgument(1, named("java.lang.reflect.Method"))),
+        "datadog.trace.instrumentation.rmi.ClientAdvice");
+  }
+}

dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/RmiServerInstrumentation.java

@@ -0,0 +1,37 @@
+package datadog.trace.instrumentation.rmi;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class RmiServerInstrumentation extends Instrumenter.Default {
+
+  public RmiServerInstrumentation() {
+    super("rmi", "rmi-server");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("java.rmi.server.RemoteServer")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod().and(isPublic()).and(not(isStatic())),
+        "datadog.trace.instrumentation.rmi.ServerAdvice");
+  }
+}

dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/ServerAdvice.java

@@ -0,0 +1,35 @@
+package datadog.trace.instrumentation.rmi;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+
+import datadog.trace.api.DDTags;
+import datadog.trace.instrumentation.api.AgentScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+
+public class ServerAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(
+      @Advice.This final Object thiz, @Advice.Origin(value = "#m") final String method) {
+    final AgentSpan span =
+        startSpan("rmi.request")
+            .setTag(DDTags.RESOURCE_NAME, thiz.getClass().getSimpleName() + "#" + method)
+            .setTag("span.origin.type", thiz.getClass().getCanonicalName());
+    return activateSpan(span, true);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void stopSpan(
+      @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+    if (scope == null) {
+      return;
+    }
+    final AgentSpan span = scope.span();
+    if (throwable != null) {
+      span.setError(true);
+      span.addThrowable(throwable);
+    }
+    scope.close();
+  }
+}

dd-java-agent/instrumentation/rmi/src/test/groovy/RmiTest.groovy

@@ -0,0 +1,174 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.utils.PortUtils
+import rmi.app.Greeter
+import rmi.app.Server
+import rmi.app.ServerLegacy
+
+import java.rmi.registry.LocateRegistry
+import java.rmi.server.UnicastRemoteObject
+
+import static datadog.trace.agent.test.asserts.ListWriterAssert.assertTraces
+import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
+class RmiTest extends AgentTestRunner {
+  def registryPort = PortUtils.randomOpenPort()
+  def serverRegistry = LocateRegistry.createRegistry(registryPort)
+  def clientRegistry = LocateRegistry.getRegistry("localhost", registryPort)
+
+  def cleanup() {
+    UnicastRemoteObject.unexportObject(serverRegistry, true)
+  }
+
+  def "Client call creates spans"() {
+    setup:
+    def server = new Server()
+    serverRegistry.rebind(Server.RMI_ID, server)
+
+    when:
+    def response = runUnderTrace("parent") {
+      def client = (Greeter) clientRegistry.lookup(Server.RMI_ID)
+      return client.hello("you")
+    }
+
+    then:
+    response.contains("Hello you")
+    assertTraces(TEST_WRITER, 2) {
+      trace(1, 2) {
+        basicSpan(it, 0, "parent")
+        span(1) {
+          resourceName "Greeter#hello"
+          operationName "rmi.invoke"
+          childOf span(0)
+          tags {
+            "span.origin.type" Greeter.canonicalName
+            defaultTags()
+          }
+        }
+      }
+      trace(0, 1) {
+        span(0) {
+          resourceName "Server#hello"
+          operationName "rmi.request"
+          tags {
+            "span.origin.type" server.class.canonicalName
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    cleanup:
+    serverRegistry.unbind("Server")
+  }
+
+  def "Calling server builtin methods doesn't create server spans"() {
+    setup:
+    def server = new Server()
+    serverRegistry.rebind(Server.RMI_ID, server)
+
+    when:
+    server.equals(new Server())
+    server.getRef()
+    server.hashCode()
+    server.toString()
+    server.getClass()
+
+    then:
+    assertTraces(TEST_WRITER, 0) {}
+
+    cleanup:
+    serverRegistry.unbind("Server")
+  }
+
+  def "Service throws exception and its propagated to spans"() {
+    setup:
+    def server = new Server()
+    serverRegistry.rebind(Server.RMI_ID, server)
+
+    when:
+    runUnderTrace("parent") {
+      def client = (Greeter) clientRegistry.lookup(Server.RMI_ID)
+      client.exceptional()
+    }
+
+    then:
+    def thrownException = thrown(RuntimeException)
+    assertTraces(TEST_WRITER, 2) {
+      trace(1, 2) {
+        basicSpan(it, 0, "parent", null, thrownException)
+        span(1) {
+          resourceName "Greeter#exceptional"
+          operationName "rmi.invoke"
+          childOf span(0)
+          errored true
+
+          tags {
+            "span.origin.type" Greeter.canonicalName
+            errorTags(RuntimeException, String)
+            defaultTags()
+          }
+        }
+      }
+      trace(0, 1) {
+        span(0) {
+          resourceName "Server#exceptional"
+          operationName "rmi.request"
+          errored true
+
+          tags {
+            "span.origin.type" server.class.canonicalName
+            errorTags(RuntimeException, String)
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    cleanup:
+    serverRegistry.unbind("Server")
+  }
+
+  def "Client call using ServerLegacy_stub creates spans"() {
+    setup:
+    def server = new ServerLegacy()
+    serverRegistry.rebind(ServerLegacy.RMI_ID, server)
+
+
+    when:
+    def response = runUnderTrace("parent") {
+      def client = (Greeter) clientRegistry.lookup(ServerLegacy.RMI_ID)
+      return client.hello("you")
+    }
+
+    then:
+    response.contains("Hello you")
+    assertTraces(TEST_WRITER, 2) {
+      trace(1, 2) {
+        basicSpan(it, 0, "parent")
+        span(1) {
+          resourceName "Greeter#hello"
+          operationName "rmi.invoke"
+          childOf span(0)
+          tags {
+            "span.origin.type" Greeter.canonicalName
+            defaultTags()
+          }
+        }
+      }
+      trace(0, 1) {
+        span(0) {
+          resourceName "ServerLegacy#hello"
+          operationName "rmi.request"
+          tags {
+            "span.origin.type" server.class.canonicalName
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    cleanup:
+    serverRegistry.unbind(ServerLegacy.RMI_ID)
+  }
+}

dd-java-agent/instrumentation/rmi/src/test/java/rmi/app/Greeter.java

@@ -0,0 +1,10 @@
+package rmi.app;
+
+import java.rmi.Remote;
+import java.rmi.RemoteException;
+
+public interface Greeter extends Remote {
+  String hello(String name) throws RemoteException;
+
+  void exceptional() throws RemoteException, RuntimeException;
+}

dd-java-agent/instrumentation/rmi/src/test/java/rmi/app/Server.java

@@ -0,0 +1,24 @@
+package rmi.app;
+
+import java.rmi.RemoteException;
+import java.rmi.server.UnicastRemoteObject;
+
+public class Server extends UnicastRemoteObject implements Greeter {
+  static String RMI_ID = Server.class.getSimpleName();
+
+  private static final long serialVersionUID = 1L;
+
+  public Server() throws RemoteException {
+    super();
+  }
+
+  @Override
+  public String hello(final String name) {
+    return "Hello " + name;
+  }
+
+  @Override
+  public void exceptional() throws RuntimeException {
+    throw new RuntimeException("expected");
+  }
+}

dd-java-agent/instrumentation/rmi/src/test/java/rmi/app/ServerLegacy.java

@@ -0,0 +1,24 @@
+package rmi.app;
+
+import java.rmi.RemoteException;
+import java.rmi.server.UnicastRemoteObject;
+
+public class ServerLegacy extends UnicastRemoteObject implements Greeter {
+  static String RMI_ID = ServerLegacy.class.getSimpleName();
+
+  private static final long serialVersionUID = 1L;
+
+  public ServerLegacy() throws RemoteException {
+    super();
+  }
+
+  @Override
+  public String hello(final String name) {
+    return "Hello " + name;
+  }
+
+  @Override
+  public void exceptional() throws RuntimeException {
+    throw new RuntimeException("expected");
+  }
+}

settings.gradle

@@ -110,8 +110,9 @@ include ':dd-java-agent:instrumentation:play-ws-2'
 include ':dd-java-agent:instrumentation:play-ws-2.1'
 include ':dd-java-agent:instrumentation:rabbitmq-amqp-2.7'
 include ':dd-java-agent:instrumentation:ratpack-1.4'
-include ':dd-java-agent:instrumentation:rxjava-1'
 include ':dd-java-agent:instrumentation:reactor-core-3.1'
+include ':dd-java-agent:instrumentation:rmi'
+include ':dd-java-agent:instrumentation:rxjava-1'
 include ':dd-java-agent:instrumentation:servlet'
 include ':dd-java-agent:instrumentation:servlet:request-2'
 include ':dd-java-agent:instrumentation:servlet:request-3'