[rmi] use ObjectEndpoint#toString() to avoid need for reflection to be able to compare object identifiers.

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

@@ -1,8 +1,6 @@
 muzzle {
   pass {
-    group = "javax.ws.rs"
+    coreJdk()
-    module = "jsr311-api"
-    versions = "[0.5,)"
   }
 }
 

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

@@ -30,7 +30,9 @@ public final class RmiClientInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".ClientDecorator", "datadog.trace.agent.decorator.BaseDecorator"
+      "datadog.trace.agent.decorator.BaseDecorator",
+      packageName + ".ClientDecorator",
+      packageName + ".ClientAdvice"
     };
   }
 

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

@@ -1,15 +0,0 @@
-package datadog.trace.instrumentation.rmi.context;
-
-import datadog.trace.bootstrap.InstrumentationContext;
-import java.rmi.server.ObjID;
-import net.bytebuddy.asm.Advice;
-
-public class ObjectEndpointConstructorAdvice {
-  @Advice.OnMethodExit(suppress = Throwable.class)
-  public static void onEnter(
-      @Advice.This final Object thiz, @Advice.Argument(value = 0) final ObjID id) {
-    if (id != null) {
-      InstrumentationContext.get(Object.class, ObjID.class).put(thiz, id);
-    }
-  }
-}

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

@@ -5,9 +5,7 @@ import static datadog.trace.instrumentation.rmi.context.StreamRemoteCallConstruc
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentSpan;
-import java.lang.reflect.Field;
 import java.rmi.Remote;
-import java.rmi.server.ObjID;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.asm.Advice;
 import sun.rmi.transport.Target;
@@ -19,9 +17,12 @@ public class ObjectTableAdvice {
   @Advice.OnMethodExit(suppress = Throwable.class)
   public static void methodExit(
       @Advice.Argument(0) final Object oe, @Advice.Return(readOnly = false) Target result) {
-    final ObjID objID = InstrumentationContext.get(Object.class, ObjID.class).get(oe);
 
-    if (!DD_CONTEXT_CALL_ID.equals(objID)) {
+    // comparing toString() output allows us to avoid using reflection to be able to compare
+    // ObjID and ObjectEndpoint objects
+    // ObjectEndpoint#toString() only returns this.objId.toString() value which is exactly
+    // what we're interested in here.
+    if (!DD_CONTEXT_CALL_ID.toString().equals(oe.toString())) {
       return;
     }
 
@@ -30,20 +31,11 @@ public class ObjectTableAdvice {
 
     result =
         new Target(
-            DUMMY_REMOTE, new ContextDispatcher(callableContextStore), DUMMY_REMOTE, objID, false);
+            DUMMY_REMOTE,
-  }
+            new ContextDispatcher(callableContextStore),
-
+            DUMMY_REMOTE,
-  public static ObjID GET_OBJ_ID(final Object oe) {
+            DD_CONTEXT_CALL_ID,
-    try {
+            false);
-      final Class<?> clazz = oe.getClass();
-      // sun.rmi.transport.ObjectEndpoint is protected and field "id" is private
-      final Field id = clazz.getDeclaredField("id");
-      id.setAccessible(true);
-      return (ObjID) id.get(oe);
-    } catch (final ReflectiveOperationException e) {
-      log.debug("Error getting object id from: {}", oe, e);
-    }
-    return null;
   }
 
   public static class DummyRemote implements Remote {}

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

@@ -2,7 +2,6 @@ package datadog.trace.instrumentation.rmi.context;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
-import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isStatic;
@@ -12,20 +11,15 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.instrumentation.api.AgentSpan;
+import java.util.Collections;
-import java.rmi.server.ObjID;
 import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import sun.rmi.transport.Connection;
-import sun.rmi.transport.ObjectTable;
-import sun.rmi.transport.StreamRemoteCall;
 
 @AutoService(Instrumenter.class)
 public class RmiContextInstrumentation extends Instrumenter.Default {
-  // TODO clean this up
 
   public RmiContextInstrumentation() {
     super("rmi", "rmi-context-propagator");
@@ -36,8 +30,8 @@ public class RmiContextInstrumentation extends Instrumenter.Default {
     return not(isInterface())
         .and(
             safeHasSuperType(
-                named(StreamRemoteCall.class.getName()) // TODO replace with string
+                named("sun.rmi.transport.StreamRemoteCall")
-                    .or(named(ObjectTable.class.getName()))
+                    .or(named("sun.rmi.transport.ObjectTable"))
                     .or(named("sun.rmi.transport.ObjectEndpoint"))));
   }
 
@@ -45,26 +39,22 @@ public class RmiContextInstrumentation extends Instrumenter.Default {
   public Map<String, String> contextStore() {
     final HashMap<String, String> contextStore = new HashMap<>();
     // thread context that stores distributed context
-    contextStore.put(Thread.class.getName(), AgentSpan.Context.class.getName());
+    contextStore.put("java.lang.Thread", "datadog.trace.instrumentation.api.AgentSpan$Context");
 
     // caching if a connection can support enhanced format
-    contextStore.put(Connection.class.getName(), Boolean.class.getName());
+    contextStore.put("sun.rmi.transport.Connection", "java.lang.Boolean");
-
-    // used to avoid reflection when instrumenting protected class ObjectEndpoint
-    contextStore.put(Object.class.getName(), ObjID.class.getName());
     return contextStore;
   }
 
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".StreamRemoteCallConstructorAdvice",
       packageName + ".ContextPayload",
       packageName + ".ContextPayload$InjectAdapter",
       packageName + ".ContextPayload$ExtractAdapter",
-      packageName + ".ObjectTableAdvice",
       packageName + ".ContextDispatcher",
-      packageName + ".ObjectEndpointConstructorAdvice",
+      packageName + ".StreamRemoteCallConstructorAdvice",
+      packageName + ".ObjectTableAdvice",
       packageName + ".ObjectTableAdvice$DummyRemote"
     };
   }
@@ -85,9 +75,6 @@ public class RmiContextInstrumentation extends Instrumenter.Default {
             .and((takesArgument(0, named("sun.rmi.transport.ObjectEndpoint")))),
         packageName + ".ObjectTableAdvice");
 
-    transformers.put(
+    return Collections.unmodifiableMap(transformers);
-        isConstructor().and(isDeclaredBy(named("sun.rmi.transport.ObjectEndpoint"))),
-        packageName + ".ObjectEndpointConstructorAdvice");
-    return transformers;
   }
 }

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

@@ -42,7 +42,9 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".ServerDecorator", "datadog.trace.agent.decorator.BaseDecorator"
+      packageName + ".ServerDecorator",
+      packageName + ".RmiServerInstrumentation$ServerAdvice",
+      "datadog.trace.agent.decorator.BaseDecorator"
     };
   }
 
@@ -55,7 +57,7 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     return singletonMap(
         isMethod().and(isPublic()).and(not(isStatic())),
-        "datadog.trace.instrumentation.rmi.server.RmiServerInstrumentation$ServerAdvice");
+        packageName + ".RmiServerInstrumentation$ServerAdvice");
   }
 
   public static class ServerAdvice {