[rmi] Split context propagation to server and client instrumentation
+ Separate extra code from advices
+ cleanup helper definitions
RENAME
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/ClientAdvice.java
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(
- DDTags.RESOURCE_NAME,
- method.getDeclaringClass().getSimpleName() + "#" + method.getName())
+ .setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method))
.setTag("span.origin.type", method.getDeclaringClass().getCanonicalName());
DECORATE.afterStart(span);

RENAME
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/ClientDecorator.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientDecorator.java
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 static final ClientDecorator DECORATE = new ClientDecorator();
+public class RmiClientDecorator extends 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";
+ }
}

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientInstrumentation.java
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",
- packageName + ".ClientAdvice"
+ "datadog.trace.agent.decorator.ClientDecorator",
+ 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");
}
}

RENAME
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/StreamRemoteCallConstructorAdvice.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.java
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 static final ObjID ACTIVATOR_ID = new ObjID(ObjID.ACTIVATOR_ID);
- public static final ObjID DGC_ID = new ObjID(ObjID.DGC_ID);
- public 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 int CONTEXT_CHECK_CALL_OP_ID = -1;
+public class ContextPropagator {
+ private static final ObjID ACTIVATOR_ID = new ObjID(ObjID.ACTIVATOR_ID);
+ private static final ObjID DGC_ID = new ObjID(ObjID.DGC_ID);
+ private static final ObjID REGISTRY_ID = new ObjID(ObjID.REGISTRY_ID);
+ public static final ObjID DD_CONTEXT_CALL_ID = new ObjID("Datadog.v1.context_call".hashCode());
+ 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(
- final ContextStore<Connection, Boolean> contextStore, final Connection c) {
- final AgentSpan span = activeSpan();
- if (span == null) {
- return;
- }
+ public static final ContextPropagator PROPAGATOR = new ContextPropagator();
+ 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 {

DELETE
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/RmiContextInstrumentation.java
-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);
- }
-}

ADD
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java
+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");
+ }
+}

ADD
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/StreamRemoteCallConstructorAdvice.java
+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);
+ }
+}

RENAME
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextDispatcher.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java
-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);

RENAME
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ObjectTableAdvice.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ObjectTableAdvice.java
-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();
+ }
}

ADD
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java
+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");
+ }
+}

ADD
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerDecorator.java
+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";
+ }
+}

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerInstrumentation.java
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.ServerDecorator.DECORATE;
+import static datadog.trace.instrumentation.rmi.server.RmiServerDecorator.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",
- packageName + ".RmiServerInstrumentation$ServerAdvice",
- "datadog.trace.agent.decorator.BaseDecorator"
+ "datadog.trace.agent.decorator.ServerDecorator",
+ "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) {
- final ContextStore<Thread, AgentSpan.Context> callableContextStore =
+ @Advice.This final Object thiz, @Advice.Origin final Method method) {
+ final ContextStore<Thread, AgentSpan.Context> threadContextStore =
InstrumentationContext.get(Thread.class, AgentSpan.Context.class);
final AgentSpan span =
- startSpan(callableContextStore)
- .setTag(DDTags.RESOURCE_NAME, thiz.getClass().getSimpleName() + "#" + method)
+ DECORATE
+ .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) {

DELETE
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/ServerDecorator.java
-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";
- }
-}

MODIFY
dd-java-agent/instrumentation/rmi/src/test/groovy/RmiTest.groovy
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)
}

[rmi] Split context propagation to server and client instrumentation

+ Separate extra code from advices
 + cleanup helper definitions