[rmi] use simple thread local for propagating context
+ add more documentation
ADD
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ThreadLocalContext.java
+package datadog.trace.bootstrap.instrumentation.rmi;
+
+import datadog.trace.instrumentation.api.AgentSpan;
+
+public class ThreadLocalContext {
+ public static final ThreadLocalContext THREAD_LOCAL_CONTEXT = new ThreadLocalContext();
+ private final ThreadLocal<AgentSpan.Context> local;
+
+ public ThreadLocalContext() {
+ local = new ThreadLocal<>();
+ }
+
+ public void set(final AgentSpan.Context context) {
+ local.set(context);
+ }
+
+ public AgentSpan.Context getAndResetContext() {
+ final AgentSpan.Context context = local.get();
+ if (context != null) {
+ local.remove();
+ }
+ return context;
+ }
+}

DELETE
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientAdvice.java
-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.RmiClientDecorator.DECORATE;
-
-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 RmiClientAdvice {
- @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, DECORATE.spanNameForMethod(method))
- .setTag("span.origin.type", method.getDeclaringClass().getCanonicalName());
-
- DECORATE.afterStart(span);
- 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;
- }
- DECORATE.onError(scope, throwable);
- scope.close();
- }
-}

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
package datadog.trace.instrumentation.rmi.client;
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.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.rmi.client.RmiClientDecorator.DECORATE;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -10,7 +14,12 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDTags;
+import datadog.trace.instrumentation.api.AgentScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import java.lang.reflect.Method;
import java.util.Map;
+import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
@@ -43,6 +52,32 @@ 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 + ".RmiClientAdvice");
+ getClass().getName() + "$RmiClientAdvice");
+ }
+
+ public static class RmiClientAdvice {
+ @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, DECORATE.spanNameForMethod(method))
+ .setTag("span.origin.type", method.getDeclaringClass().getCanonicalName());
+
+ DECORATE.afterStart(span);
+ 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;
+ }
+ DECORATE.onError(scope, throwable);
+ scope.close();
+ }
}
}

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java
package datadog.trace.instrumentation.rmi.context;
+import static datadog.trace.instrumentation.api.AgentTracer.propagate;
+
import datadog.trace.instrumentation.api.AgentPropagation;
+import datadog.trace.instrumentation.api.AgentSpan;
import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
-import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
+/** ContextPayload wraps context information shared between client and server */
@Slf4j
-public class ContextPayload implements Serializable {
+public class ContextPayload {
@Getter private final Map<String, String> context;
public static final ExtractAdapter GETTER = new ExtractAdapter();
public static final InjectAdapter SETTER = new InjectAdapter();
@@ -24,6 +27,12 @@ public class ContextPayload implements Serializable {
this.context = context;
}
+ public static ContextPayload from(final AgentSpan span) {
+ final ContextPayload payload = new ContextPayload();
+ propagate().inject(span, payload, SETTER);
+ return payload;
+ }
+
public static ContextPayload read(final ObjectInput oi) throws IOException {
try {
final Object object = oi.readObject();

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.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.propagate;
-import static datadog.trace.instrumentation.rmi.context.ContextPayload.SETTER;
-
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.instrumentation.api.AgentSpan;
import java.io.IOException;
@@ -16,44 +13,55 @@ import sun.rmi.transport.TransportConstants;
@Slf4j
public class ContextPropagator {
+ // Internal RMI object ids that we don't want to trace
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);
+
+ // RMI object id used to identify DataDog instrumentation
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;
- public static boolean isRMIInternalObject(final ObjID id) {
+ // Operation id used for checking context propagation is possible
+ // RMI expects these operations to have negative identifier, as positive ones mean legacy
+ // precompiled Stubs would be used instead
+ private static final int CONTEXT_CHECK_CALL_OPERATION_ID = -1;
+ // Seconds step of context propagation which contains actual payload
+ private static final int CONTEXT_PAYLOAD_OPERATION_ID = -2;
+
+ public static final ContextPropagator PROPAGATOR = new ContextPropagator();
+
+ public boolean isRMIInternalObject(final ObjID id) {
return ACTIVATOR_ID.equals(id) || DGC_ID.equals(id) || REGISTRY_ID.equals(id);
}
- public static final ContextPropagator PROPAGATOR = new ContextPropagator();
+ public boolean isOperationWithPayload(final int operationId) {
+ return operationId == CONTEXT_PAYLOAD_OPERATION_ID;
+ }
public void attemptToPropagateContext(
- final ContextStore<Connection, Boolean> contextStore,
+ final ContextStore<Connection, Boolean> knownConnections,
final Connection c,
final AgentSpan span) {
- if (checkIfContextCanBePassed(contextStore, c)) {
- final ContextPayload payload = new ContextPayload();
- propagate().inject(span, payload, SETTER);
- if (!syntheticCall(c, payload, CONTEXT_PASS_OPERATION_ID)) {
- log.debug("Couldn't propagate context");
+ if (checkIfContextCanBePassed(knownConnections, c)) {
+ if (!syntheticCall(c, ContextPayload.from(span), CONTEXT_PAYLOAD_OPERATION_ID)) {
+ log.debug("Couldn't send context payload");
}
}
}
private boolean checkIfContextCanBePassed(
- final ContextStore<Connection, Boolean> contextStore, final Connection c) {
- final Boolean storedResult = contextStore.get(c);
+ final ContextStore<Connection, Boolean> knownConnections, final Connection c) {
+ final Boolean storedResult = knownConnections.get(c);
if (storedResult != null) {
return storedResult;
}
- final boolean result = syntheticCall(c, null, CONTEXT_CHECK_CALL_OP_ID);
- contextStore.put(c, result);
+ final boolean result = syntheticCall(c, null, CONTEXT_CHECK_CALL_OPERATION_ID);
+ knownConnections.put(c, result);
return result;
}
+ /** @returns true when no error happened during call */
private boolean syntheticCall(
final Connection c, final ContextPayload payload, final int operationId) {
final StreamRemoteCall shareContextCall = new StreamRemoteCall(c);
@@ -68,9 +76,10 @@ public class ContextPropagator {
out.writeInt(operationId); // in normal call this is method number (operation index)
out.writeLong(operationId); // in normal RMI call this holds stub/skeleton hash
- // if method is not found by uninstrumented code then writing payload will cause an exception
- // in
- // RMI server - as the payload will be interpreted as another call
+ // Payload should be sent only after we make sure we're connected to instrumented server
+ //
+ // if method is not found by un-instrumented code then writing payload will cause an exception
+ // in RMI server - as the payload will be interpreted as another call
// but it will not be parsed correctly - closing connection
if (payload != null) {
payload.write(out);

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java
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 datadog.trace.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.rmi.context.ContextPropagator.PROPAGATOR;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
@@ -10,11 +12,38 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.instrumentation.api.AgentSpan;
+import java.rmi.server.ObjID;
import java.util.Map;
+import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
+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
+ */
@AutoService(Instrumenter.class)
public class RmiClientContextInstrumentation extends Instrumenter.Default {
@@ -49,6 +78,29 @@ public class RmiClientContextInstrumentation extends Instrumenter.Default {
isConstructor()
.and(takesArgument(0, named("sun.rmi.transport.Connection")))
.and(takesArgument(1, named("java.rmi.server.ObjID"))),
- packageName + ".StreamRemoteCallConstructorAdvice");
+ getClass().getName() + "$StreamRemoteCallConstructorAdvice");
+ }
+
+ public static 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 (PROPAGATOR.isRMIInternalObject(id)) {
+ return;
+ }
+ final AgentSpan activeSpan = activeSpan();
+ if (activeSpan == null) {
+ return;
+ }
+
+ final ContextStore<Connection, Boolean> knownConnections =
+ InstrumentationContext.get(Connection.class, Boolean.class);
+
+ PROPAGATOR.attemptToPropagateContext(knownConnections, c, activeSpan);
+ }
}
}

DELETE
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);
- }
-}

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java
package datadog.trace.instrumentation.rmi.context.server;
+import static datadog.trace.bootstrap.instrumentation.rmi.ThreadLocalContext.THREAD_LOCAL_CONTEXT;
import static datadog.trace.instrumentation.api.AgentTracer.propagate;
import static datadog.trace.instrumentation.rmi.context.ContextPayload.GETTER;
+import static datadog.trace.instrumentation.rmi.context.ContextPropagator.DD_CONTEXT_CALL_ID;
+import static datadog.trace.instrumentation.rmi.context.ContextPropagator.PROPAGATOR;
-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;
import java.rmi.server.RemoteCall;
import sun.rmi.server.Dispatcher;
-
+import sun.rmi.transport.Target;
+
+/**
+ * ContextDispatcher is responsible for handling both initial context propagation check call and
+ * following call which carries payload
+ *
+ * <p>Context propagation check is only expected not to throw any exception, hinting to the client
+ * that its communicating with an instrumented server. Non instrumented server would've thrown
+ * UnknownObjectException
+ *
+ * <p>Because caching of the result after first call on a connection, only payload calls are
+ * expected
+ */
public class ContextDispatcher implements Dispatcher {
+ private static final ContextDispatcher CONTEXT_DISPATCHER = new ContextDispatcher();
+ private static final NoopRemote NOOP_REMOTE = new NoopRemote();
- private final ContextStore<Thread, AgentSpan.Context> callableContextStore;
-
- public ContextDispatcher(final ContextStore<Thread, AgentSpan.Context> callableContextStore) {
- this.callableContextStore = callableContextStore;
+ public static Target newDispatcherTarget() {
+ return new Target(NOOP_REMOTE, CONTEXT_DISPATCHER, NOOP_REMOTE, DD_CONTEXT_CALL_ID, false);
}
@Override
@@ -27,12 +40,11 @@ public class ContextDispatcher implements Dispatcher {
final int operationId = in.readInt();
in.readLong(); // skip 8 bytes
- if (operationId == ContextPropagator.CONTEXT_PASS_OPERATION_ID) {
+ if (PROPAGATOR.isOperationWithPayload(operationId)) {
final ContextPayload payload = ContextPayload.read(in);
if (payload != null) {
final AgentSpan.Context context = propagate().extract(payload, GETTER);
-
- callableContextStore.put(Thread.currentThread(), context);
+ THREAD_LOCAL_CONTEXT.set(context);
}
}
@@ -44,4 +56,6 @@ public class ContextDispatcher implements Dispatcher {
call.releaseOutputStream();
call.done();
}
+
+ public static class NoopRemote implements Remote {}
}

DELETE
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ObjectTableAdvice.java
-package datadog.trace.instrumentation.rmi.context.server;
-
-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 net.bytebuddy.asm.Advice;
-import sun.rmi.transport.Target;
-
-public class ObjectTableAdvice {
-
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void methodExit(
- @Advice.Argument(0) final Object oe, @Advice.Return(readOnly = false) Target result) {
-
- // 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;
- }
-
- final ContextStore<Thread, AgentSpan.Context> callableContextStore =
- InstrumentationContext.get(Thread.class, AgentSpan.Context.class);
-
- result =
- new Target(
- NoopRemote.NOOP_REMOTE,
- new ContextDispatcher(callableContextStore),
- NoopRemote.NOOP_REMOTE,
- DD_CONTEXT_CALL_ID,
- false);
- }
-
- public static class NoopRemote implements Remote {
- public static final NoopRemote NOOP_REMOTE = new NoopRemote();
- }
-}

MODIFY
dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java
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 datadog.trace.instrumentation.rmi.context.ContextPropagator.DD_CONTEXT_CALL_ID;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -12,9 +13,11 @@ 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.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
+import sun.rmi.transport.Target;
@AutoService(Instrumenter.class)
public class RmiServerContextInstrumentation extends Instrumenter.Default {
@@ -28,11 +31,6 @@ public class RmiServerContextInstrumentation extends Instrumenter.Default {
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[] {
@@ -41,7 +39,7 @@ public class RmiServerContextInstrumentation extends Instrumenter.Default {
"datadog.trace.instrumentation.rmi.context.ContextPayload",
"datadog.trace.instrumentation.rmi.context.ContextPropagator",
packageName + ".ContextDispatcher",
- packageName + ".ObjectTableAdvice$NoopRemote"
+ packageName + ".ContextDispatcher$NoopRemote"
};
}
@@ -52,6 +50,22 @@ public class RmiServerContextInstrumentation extends Instrumenter.Default {
.and(isStatic())
.and(named("getTarget"))
.and((takesArgument(0, named("sun.rmi.transport.ObjectEndpoint")))),
- packageName + ".ObjectTableAdvice");
+ getClass().getName() + "$ObjectTableAdvice");
+ }
+
+ public static class ObjectTableAdvice {
+
+ @Advice.OnMethodExit(suppress = Throwable.class)
+ public static void methodExit(
+ @Advice.Argument(0) final Object oe, @Advice.Return(readOnly = false) Target result) {
+ // 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;
+ }
+ result = ContextDispatcher.newDispatcherTarget();
+ }
}
}

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

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
package datadog.trace.instrumentation.rmi.server;
import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.bootstrap.instrumentation.rmi.ThreadLocalContext.THREAD_LOCAL_CONTEXT;
import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.instrumentation.rmi.server.RmiServerDecorator.DECORATE;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
@@ -14,12 +16,9 @@ import static net.bytebuddy.matcher.ElementMatchers.not;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.api.DDTags;
-import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.instrumentation.api.AgentScope;
import datadog.trace.instrumentation.api.AgentSpan;
import java.lang.reflect.Method;
-import java.util.Collections;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
@@ -33,11 +32,6 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
super("rmi", "rmi-server");
}
- @Override
- public Map<String, String> contextStore() {
- return Collections.singletonMap(Thread.class.getName(), AgentSpan.Context.class.getName());
- }
-
@Override
public String[] helperClassNames() {
return new String[] {
@@ -55,21 +49,23 @@ public final class RmiServerInstrumentation extends Instrumenter.Default {
@Override
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
return singletonMap(
- isMethod().and(isPublic()).and(not(isStatic())),
- packageName + ".RmiServerInstrumentation$ServerAdvice");
+ isMethod().and(isPublic()).and(not(isStatic())), getClass().getName() + "$ServerAdvice");
}
public static class ServerAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class, inline = true)
public static AgentScope onEnter(
@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.Context context = THREAD_LOCAL_CONTEXT.getAndResetContext();
+
+ final AgentSpan span;
+ if (context == null) {
+ span = startSpan("rmi.request");
+ } else {
+ span = startSpan("rmi.request", context);
+ }
- final AgentSpan span =
- DECORATE
- .startSpanWithContext(threadContextStore)
- .setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method))
+ span.setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method))
.setTag("span.origin.type", thiz.getClass().getCanonicalName());
DECORATE.afterStart(span);

[rmi] use simple thread local for propagating context

+ add more documentation