[rmi] Split context propagation to server and client instrumentation

+ Separate extra code from advices
 + cleanup helper definitions

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

@@ -3,7 +3,7 @@ package datadog.trace.instrumentation.rmi.client;
 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 static datadog.trace.instrumentation.rmi.client.ClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.rmi.client.RmiClientDecorator.DECORATE;
 
 import datadog.trace.api.DDTags;
 import datadog.trace.instrumentation.api.AgentScope;
@@ -11,7 +11,7 @@ import datadog.trace.instrumentation.api.AgentSpan;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
 
-public class ClientAdvice {
+public class RmiClientAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static AgentScope onEnter(@Advice.Argument(value = 1) final Method method) {
     if (activeSpan() == null) {
@@ -19,9 +19,7 @@ public class ClientAdvice {
     }
     final AgentSpan span =
         startSpan("rmi.invoke")
-            .setTag(
+            .setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method))
-                DDTags.RESOURCE_NAME,
-                method.getDeclaringClass().getSimpleName() + "#" + method.getName())
             .setTag("span.origin.type", method.getDeclaringClass().getCanonicalName());
 
     DECORATE.afterStart(span);

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

@@ -1,10 +1,10 @@
 package datadog.trace.instrumentation.rmi.client;
 
-import datadog.trace.agent.decorator.BaseDecorator;
+import datadog.trace.agent.decorator.ClientDecorator;
 import datadog.trace.api.DDSpanTypes;
 
-public class ClientDecorator extends BaseDecorator {
+public class RmiClientDecorator extends ClientDecorator {
-  public static final ClientDecorator DECORATE = new ClientDecorator();
+  public static final RmiClientDecorator DECORATE = new RmiClientDecorator();
 
   @Override
   protected String[] instrumentationNames() {
@@ -20,4 +20,9 @@ public class ClientDecorator extends BaseDecorator {
   protected String component() {
     return "rmi-client";
   }
+
+  @Override
+  protected String service() {
+    return "rmi";
+  }
 }

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

@@ -31,8 +31,8 @@ public final class RmiClientInstrumentation extends Instrumenter.Default {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.trace.agent.decorator.BaseDecorator",
-      packageName + ".ClientDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
-      packageName + ".ClientAdvice"
+      packageName + ".RmiClientDecorator"
     };
   }
 
@@ -43,6 +43,6 @@ public final class RmiClientInstrumentation extends Instrumenter.Default {
             .and(named("invoke"))
             .and(takesArgument(0, named("java.rmi.Remote")))
             .and(takesArgument(1, named("java.lang.reflect.Method"))),
-        packageName + ".ClientAdvice");
+        packageName + ".RmiClientAdvice");
   }
 }

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

@@ -1,78 +1,38 @@
 package datadog.trace.instrumentation.rmi.context;
 
-import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.instrumentation.rmi.context.ContextPayload.SETTER;
 
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentSpan;
 import java.io.IOException;
 import java.io.ObjectOutput;
 import java.rmi.NoSuchObjectException;
 import java.rmi.server.ObjID;
 import lombok.extern.slf4j.Slf4j;
-import net.bytebuddy.asm.Advice;
 import sun.rmi.transport.Connection;
 import sun.rmi.transport.StreamRemoteCall;
 import sun.rmi.transport.TransportConstants;
 
-/**
- * Main entry point for transferring context between RMI service.
- *
- * <p>It injects into StreamRemoteCall constructor used for invoking remote tasks and performs a
- * backwards compatible check to ensure if the other side is prepared to receive context propagation
- * messages then if successful sends a context propagation message
- *
- * <p>Context propagation consist of a Serialized HashMap with all data set by usual context
- * injection, which includes things like sampling priority, trace and parent id
- *
- * <p>As well as optional baggage items
- *
- * <p>On the other side of the communication a special Dispatcher is created when a message with
- * DD_CONTEXT_CALL_ID is received.
- *
- * <p>If the server is not instrumented first call will gracefully fail just like any other unknown
- * call. With small caveat that this first call needs to *not* have any parameters, since those will
- * not be read from connection and instead will be interpreted as another remote instruction, but
- * that instruction will essentially be garbage data and will cause the parsing loop to throw
- * exception and shutdown the connection which we do not want
- */
 @Slf4j
-public class StreamRemoteCallConstructorAdvice {
+public class ContextPropagator {
-  public static final ObjID ACTIVATOR_ID = new ObjID(ObjID.ACTIVATOR_ID);
+  private static final ObjID ACTIVATOR_ID = new ObjID(ObjID.ACTIVATOR_ID);
-  public static final ObjID DGC_ID = new ObjID(ObjID.DGC_ID);
+  private static final ObjID DGC_ID = new ObjID(ObjID.DGC_ID);
-  public static final ObjID REGISTRY_ID = new ObjID(ObjID.REGISTRY_ID);
+  private static final ObjID REGISTRY_ID = new ObjID(ObjID.REGISTRY_ID);
-  public static final ObjID DD_CONTEXT_CALL_ID = new ObjID("Datadog.context_call.v2".hashCode());
+  public static final ObjID DD_CONTEXT_CALL_ID = new ObjID("Datadog.v1.context_call".hashCode());
-  public static final int CONTEXT_CHECK_CALL_OP_ID = -1;
+  private static final int CONTEXT_CHECK_CALL_OP_ID = -1;
   public static final int CONTEXT_PASS_OPERATION_ID = -2;
 
-  @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static void onEnter(
-      @Advice.Argument(value = 0) final Connection c, @Advice.Argument(value = 1) final ObjID id) {
-    if (!c.isReusable()) {
-      return;
-    }
-    if (isRMIInternalObject(id)) {
-      return;
-    }
-
-    final ContextStore<Connection, Boolean> contextStore =
-        InstrumentationContext.get(Connection.class, Boolean.class);
-    attemptToPropagateContext(contextStore, c);
-  }
-
   public static boolean isRMIInternalObject(final ObjID id) {
     return ACTIVATOR_ID.equals(id) || DGC_ID.equals(id) || REGISTRY_ID.equals(id);
   }
 
-  public static void attemptToPropagateContext(
+  public static final ContextPropagator PROPAGATOR = new ContextPropagator();
-      final ContextStore<Connection, Boolean> contextStore, final Connection c) {
-    final AgentSpan span = activeSpan();
-    if (span == null) {
-      return;
-    }
 
+  public void attemptToPropagateContext(
+      final ContextStore<Connection, Boolean> contextStore,
+      final Connection c,
+      final AgentSpan span) {
     if (checkIfContextCanBePassed(contextStore, c)) {
       final ContextPayload payload = new ContextPayload();
       propagate().inject(span, payload, SETTER);
@@ -82,7 +42,7 @@ public class StreamRemoteCallConstructorAdvice {
     }
   }
 
-  private static boolean checkIfContextCanBePassed(
+  private boolean checkIfContextCanBePassed(
       final ContextStore<Connection, Boolean> contextStore, final Connection c) {
     final Boolean storedResult = contextStore.get(c);
     if (storedResult != null) {
@@ -94,7 +54,7 @@ public class StreamRemoteCallConstructorAdvice {
     return result;
   }
 
-  private static boolean syntheticCall(
+  private boolean syntheticCall(
       final Connection c, final ContextPayload payload, final int operationId) {
     final StreamRemoteCall shareContextCall = new StreamRemoteCall(c);
     try {

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

@@ -1,80 +0,0 @@
-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.isInterface;
-import static net.bytebuddy.matcher.ElementMatchers.isMethod;
-import static net.bytebuddy.matcher.ElementMatchers.isStatic;
-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.Collections;
-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;
-
-@AutoService(Instrumenter.class)
-public class RmiContextInstrumentation extends Instrumenter.Default {
-
-  public RmiContextInstrumentation() {
-    super("rmi", "rmi-context-propagator");
-  }
-
-  @Override
-  public ElementMatcher<? super TypeDescription> typeMatcher() {
-    return not(isInterface())
-        .and(
-            safeHasSuperType(
-                named("sun.rmi.transport.StreamRemoteCall")
-                    .or(named("sun.rmi.transport.ObjectTable"))
-                    .or(named("sun.rmi.transport.ObjectEndpoint"))));
-  }
-
-  @Override
-  public Map<String, String> contextStore() {
-    final HashMap<String, String> contextStore = new HashMap<>();
-    // thread context that stores distributed context
-    contextStore.put("java.lang.Thread", "datadog.trace.instrumentation.api.AgentSpan$Context");
-
-    // caching if a connection can support enhanced format
-    contextStore.put("sun.rmi.transport.Connection", "java.lang.Boolean");
-    return contextStore;
-  }
-
-  @Override
-  public String[] helperClassNames() {
-    return new String[] {
-      packageName + ".ContextPayload",
-      packageName + ".ContextPayload$InjectAdapter",
-      packageName + ".ContextPayload$ExtractAdapter",
-      packageName + ".ContextDispatcher",
-      packageName + ".StreamRemoteCallConstructorAdvice",
-      packageName + ".ObjectTableAdvice",
-      packageName + ".ObjectTableAdvice$DummyRemote"
-    };
-  }
-
-  @Override
-  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
-    transformers.put(
-        isConstructor()
-            .and(takesArgument(0, named("sun.rmi.transport.Connection")))
-            .and(takesArgument(1, named("java.rmi.server.ObjID"))),
-        packageName + ".StreamRemoteCallConstructorAdvice");
-
-    transformers.put(
-        isMethod()
-            .and(isStatic())
-            .and(named("getTarget"))
-            .and((takesArgument(0, named("sun.rmi.transport.ObjectEndpoint")))),
-        packageName + ".ObjectTableAdvice");
-
-    return Collections.unmodifiableMap(transformers);
-  }
-}

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

@@ -0,0 +1,54 @@
+package datadog.trace.instrumentation.rmi.context.client;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+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 class RmiClientContextInstrumentation extends Instrumenter.Default {
+
+  public RmiClientContextInstrumentation() {
+    super("rmi", "rmi-context-propagator", "rmi-client-context-propagator");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("sun.rmi.transport.StreamRemoteCall")));
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    // caching if a connection can support enhanced format
+    return singletonMap("sun.rmi.transport.Connection", "java.lang.Boolean");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.rmi.context.ContextPayload$InjectAdapter",
+      "datadog.trace.instrumentation.rmi.context.ContextPayload$ExtractAdapter",
+      "datadog.trace.instrumentation.rmi.context.ContextPayload",
+      "datadog.trace.instrumentation.rmi.context.ContextPropagator"
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isConstructor()
+            .and(takesArgument(0, named("sun.rmi.transport.Connection")))
+            .and(takesArgument(1, named("java.rmi.server.ObjID"))),
+        packageName + ".StreamRemoteCallConstructorAdvice");
+  }
+}

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

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.rmi.context.client;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.rmi.context.ContextPropagator.PROPAGATOR;
+
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.instrumentation.api.AgentSpan;
+import datadog.trace.instrumentation.rmi.context.ContextPropagator;
+import java.rmi.server.ObjID;
+import net.bytebuddy.asm.Advice;
+import sun.rmi.transport.Connection;
+
+/**
+ * Main entry point for transferring context between RMI service.
+ *
+ * <p>It injects into StreamRemoteCall constructor used for invoking remote tasks and performs a
+ * backwards compatible check to ensure if the other side is prepared to receive context propagation
+ * messages then if successful sends a context propagation message
+ *
+ * <p>Context propagation consist of a Serialized HashMap with all data set by usual context
+ * injection, which includes things like sampling priority, trace and parent id
+ *
+ * <p>As well as optional baggage items
+ *
+ * <p>On the other side of the communication a special Dispatcher is created when a message with
+ * DD_CONTEXT_CALL_ID is received.
+ *
+ * <p>If the server is not instrumented first call will gracefully fail just like any other unknown
+ * call. With small caveat that this first call needs to *not* have any parameters, since those will
+ * not be read from connection and instead will be interpreted as another remote instruction, but
+ * that instruction will essentially be garbage data and will cause the parsing loop to throw
+ * exception and shutdown the connection which we do not want
+ */
+public class StreamRemoteCallConstructorAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void onEnter(
+      @Advice.Argument(value = 0) final Connection c, @Advice.Argument(value = 1) final ObjID id) {
+    if (!c.isReusable()) {
+      return;
+    }
+    if (ContextPropagator.isRMIInternalObject(id)) {
+      return;
+    }
+    final AgentSpan activeSpan = activeSpan();
+    if (activeSpan == null) {
+      return;
+    }
+
+    final ContextStore<Connection, Boolean> contextStore =
+        InstrumentationContext.get(Connection.class, Boolean.class);
+
+    PROPAGATOR.attemptToPropagateContext(contextStore, c, activeSpan);
+  }
+}

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

@@ -1,10 +1,12 @@
-package datadog.trace.instrumentation.rmi.context;
+package datadog.trace.instrumentation.rmi.context.server;
 
 import static datadog.trace.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.instrumentation.rmi.context.ContextPayload.GETTER;
 
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.instrumentation.api.AgentSpan;
+import datadog.trace.instrumentation.rmi.context.ContextPayload;
+import datadog.trace.instrumentation.rmi.context.ContextPropagator;
 import java.io.IOException;
 import java.io.ObjectInput;
 import java.rmi.Remote;
@@ -25,7 +27,7 @@ public class ContextDispatcher implements Dispatcher {
     final int operationId = in.readInt();
     in.readLong(); // skip 8 bytes
 
-    if (operationId == StreamRemoteCallConstructorAdvice.CONTEXT_PASS_OPERATION_ID) {
+    if (operationId == ContextPropagator.CONTEXT_PASS_OPERATION_ID) {
       final ContextPayload payload = ContextPayload.read(in);
       if (payload != null) {
         final AgentSpan.Context context = propagate().extract(payload, GETTER);

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

@@ -1,18 +1,15 @@
-package datadog.trace.instrumentation.rmi.context;
+package datadog.trace.instrumentation.rmi.context.server;
 
-import static datadog.trace.instrumentation.rmi.context.StreamRemoteCallConstructorAdvice.DD_CONTEXT_CALL_ID;
+import static datadog.trace.instrumentation.rmi.context.ContextPropagator.DD_CONTEXT_CALL_ID;
 
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentSpan;
 import java.rmi.Remote;
-import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.asm.Advice;
 import sun.rmi.transport.Target;
 
-@Slf4j
 public class ObjectTableAdvice {
-  public static final DummyRemote DUMMY_REMOTE = new DummyRemote();
 
   @Advice.OnMethodExit(suppress = Throwable.class)
   public static void methodExit(
@@ -31,12 +28,14 @@ public class ObjectTableAdvice {
 
     result =
         new Target(
-            DUMMY_REMOTE,
+            NoopRemote.NOOP_REMOTE,
             new ContextDispatcher(callableContextStore),
-            DUMMY_REMOTE,
+            NoopRemote.NOOP_REMOTE,
             DD_CONTEXT_CALL_ID,
             false);
   }
 
-  public static class DummyRemote implements Remote {}
+  public static class NoopRemote implements Remote {
+    public static final NoopRemote NOOP_REMOTE = new NoopRemote();
+  }
 }

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

@@ -0,0 +1,57 @@
+package datadog.trace.instrumentation.rmi.context.server;
+
+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.isStatic;
+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 class RmiServerContextInstrumentation extends Instrumenter.Default {
+
+  public RmiServerContextInstrumentation() {
+    super("rmi", "rmi-context-propagator", "rmi-server-context-propagator");
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("sun.rmi.transport.ObjectTable")));
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("java.lang.Thread", "datadog.trace.instrumentation.api.AgentSpan$Context");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.rmi.context.ContextPayload$InjectAdapter",
+      "datadog.trace.instrumentation.rmi.context.ContextPayload$ExtractAdapter",
+      "datadog.trace.instrumentation.rmi.context.ContextPayload",
+      "datadog.trace.instrumentation.rmi.context.ContextPropagator",
+      packageName + ".ContextDispatcher",
+      packageName + ".ObjectTableAdvice$NoopRemote"
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(isStatic())
+            .and(named("getTarget"))
+            .and((takesArgument(0, named("sun.rmi.transport.ObjectEndpoint")))),
+        packageName + ".ObjectTableAdvice");
+  }
+}

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

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.rmi.server;
+
+import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+
+import datadog.trace.agent.decorator.ServerDecorator;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.instrumentation.api.AgentSpan;
+
+public class RmiServerDecorator extends ServerDecorator {
+  public static final RmiServerDecorator DECORATE = new RmiServerDecorator();
+
+  public AgentSpan startSpanWithContext(
+      final ContextStore<Thread, AgentSpan.Context> contextStore) {
+    if (activeSpan() != null) {
+      return startSpan("rmi.request");
+    }
+
+    final AgentSpan.Context context = contextStore.get(Thread.currentThread());
+
+    if (context == null) {
+      return startSpan("rmi.request");
+    } else {
+      return startSpan("rmi.request", context);
+    }
+  }
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"rmi"};
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.RPC;
+  }
+
+  @Override
+  protected String component() {
+    return "rmi-server";
+  }
+}

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

@@ -2,8 +2,7 @@ package datadog.trace.instrumentation.rmi.server;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
 import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.rmi.server.RmiServerDecorator.DECORATE;
-import static datadog.trace.instrumentation.rmi.server.ServerDecorator.DECORATE;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -19,7 +18,7 @@ import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.api.AgentSpan;
-import datadog.trace.instrumentation.api.AgentTracer;
+import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -42,9 +41,9 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".ServerDecorator",
+      "datadog.trace.agent.decorator.ServerDecorator",
-      packageName + ".RmiServerInstrumentation$ServerAdvice",
+      "datadog.trace.agent.decorator.BaseDecorator",
-      "datadog.trace.agent.decorator.BaseDecorator"
+      packageName + ".RmiServerDecorator"
     };
   }
 
@@ -63,33 +62,20 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
   public static class ServerAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class, inline = true)
     public static AgentScope onEnter(
-        @Advice.This final Object thiz, @Advice.Origin(value = "#m") final String method) {
+        @Advice.This final Object thiz, @Advice.Origin final Method method) {
-      final ContextStore<Thread, AgentSpan.Context> callableContextStore =
+      final ContextStore<Thread, AgentSpan.Context> threadContextStore =
           InstrumentationContext.get(Thread.class, AgentSpan.Context.class);
 
       final AgentSpan span =
-          startSpan(callableContextStore)
+          DECORATE
-              .setTag(DDTags.RESOURCE_NAME, thiz.getClass().getSimpleName() + "#" + method)
+              .startSpanWithContext(threadContextStore)
+              .setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method))
               .setTag("span.origin.type", thiz.getClass().getCanonicalName());
+
       DECORATE.afterStart(span);
       return activateSpan(span, true);
     }
 
-    public static AgentSpan startSpan(
-        final ContextStore<Thread, AgentSpan.Context> callableContextStore) {
-      if (activeSpan() != null) {
-        return AgentTracer.startSpan("rmi.request");
-      }
-
-      final AgentSpan.Context context = callableContextStore.get(Thread.currentThread());
-
-      if (context == null) {
-        return AgentTracer.startSpan("rmi.request");
-      } else {
-        return AgentTracer.startSpan("rmi.request", context);
-      }
-    }
-
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {

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

@@ -1,23 +0,0 @@
-package datadog.trace.instrumentation.rmi.server;
-
-import datadog.trace.agent.decorator.BaseDecorator;
-import datadog.trace.api.DDSpanTypes;
-
-public class ServerDecorator extends BaseDecorator {
-  public static final ServerDecorator DECORATE = new ServerDecorator();
-
-  @Override
-  protected String[] instrumentationNames() {
-    return new String[] {"rmi"};
-  }
-
-  @Override
-  protected String spanType() {
-    return DDSpanTypes.RPC;
-  }
-
-  @Override
-  protected String component() {
-    return "rmi-server";
-  }
-}

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

@@ -39,12 +39,13 @@ class RmiTest extends AgentTestRunner {
       trace(1, 2) {
         basicSpan(it, 0, "parent")
         span(1) {
-          resourceName "Greeter#hello"
+          resourceName "Greeter.hello"
           operationName "rmi.invoke"
           childOf span(0)
           spanType DDSpanTypes.RPC
 
           tags {
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
             "$Tags.COMPONENT" "rmi-client"
             "span.origin.type" Greeter.canonicalName
             defaultTags()
@@ -53,21 +54,23 @@ class RmiTest extends AgentTestRunner {
       }
       trace(0, 2) {
         span(0) {
-          resourceName "Server#hello"
+          resourceName "Server.hello"
           operationName "rmi.request"
           spanType DDSpanTypes.RPC
 
           tags {
-            "span.origin.type" server.class.canonicalName
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             "$Tags.COMPONENT" "rmi-server"
+            "span.origin.type" server.class.canonicalName
             defaultTags(true)
           }
         }
         span(1) {
-          resourceName "Server#someMethod"
+          resourceName "Server.someMethod"
           operationName "rmi.request"
           spanType DDSpanTypes.RPC
           tags {
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             "$Tags.COMPONENT" "rmi-server"
             "span.origin.type" server.class.canonicalName
             defaultTags()
@@ -116,13 +119,14 @@ class RmiTest extends AgentTestRunner {
       trace(1, 2) {
         basicSpan(it, 0, "parent", null, thrownException)
         span(1) {
-          resourceName "Greeter#exceptional"
+          resourceName "Greeter.exceptional"
           operationName "rmi.invoke"
           childOf span(0)
           errored true
           spanType DDSpanTypes.RPC
 
           tags {
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
             "$Tags.COMPONENT" "rmi-client"
             "span.origin.type" Greeter.canonicalName
             errorTags(RuntimeException, String)
@@ -132,12 +136,13 @@ class RmiTest extends AgentTestRunner {
       }
       trace(0, 1) {
         span(0) {
-          resourceName "Server#exceptional"
+          resourceName "Server.exceptional"
           operationName "rmi.request"
           errored true
           spanType DDSpanTypes.RPC
 
           tags {
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             "span.origin.type" server.class.canonicalName
             "$Tags.COMPONENT" "rmi-server"
             errorTags(RuntimeException, String)
@@ -169,11 +174,12 @@ class RmiTest extends AgentTestRunner {
       trace(1, 2) {
         basicSpan(it, 0, "parent")
         span(1) {
-          resourceName "Greeter#hello"
+          resourceName "Greeter.hello"
           operationName "rmi.invoke"
           spanType DDSpanTypes.RPC
           childOf span(0)
           tags {
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
             "$Tags.COMPONENT" "rmi-client"
             "span.origin.type" Greeter.canonicalName
             defaultTags()
@@ -184,11 +190,12 @@ class RmiTest extends AgentTestRunner {
       trace(0, 1) {
         span(0) {
           childOf parentSpan
-          resourceName "ServerLegacy#hello"
+          resourceName "ServerLegacy.hello"
           operationName "rmi.request"
           spanType DDSpanTypes.RPC
           tags {
             "$Tags.COMPONENT" "rmi-server"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             "span.origin.type" server.class.canonicalName
             defaultTags(true)
           }