diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ThreadLocalContext.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ThreadLocalContext.java new file mode 100644 index 0000000000..d4c7e6fc81 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rmi/ThreadLocalContext.java @@ -0,0 +1,24 @@ +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 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; + } +} diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientAdvice.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientAdvice.java deleted file mode 100644 index 2819807937..0000000000 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientAdvice.java +++ /dev/null @@ -1,38 +0,0 @@ -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(); - } -} diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientInstrumentation.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientInstrumentation.java index 76a8e92a08..259372ed30 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientInstrumentation.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/client/RmiClientInstrumentation.java @@ -1,6 +1,10 @@ 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(); + } } } diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java index 25de0574c2..a7fe9b1ff5 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPayload.java @@ -1,17 +1,20 @@ 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 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(); diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.java index 7b8aae8548..2c8a829f7d 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/ContextPropagator.java @@ -1,8 +1,5 @@ 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); - 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) { - return ACTIVATOR_ID.equals(id) || DGC_ID.equals(id) || REGISTRY_ID.equals(id); - } + // RMI object id used to identify DataDog instrumentation + public static final ObjID DD_CONTEXT_CALL_ID = new ObjID("Datadog.v1.context_call".hashCode()); + + // 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 boolean isOperationWithPayload(final int operationId) { + return operationId == CONTEXT_PAYLOAD_OPERATION_ID; + } + public void attemptToPropagateContext( - final ContextStore contextStore, + final ContextStore 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 contextStore, final Connection c) { - final Boolean storedResult = contextStore.get(c); + final ContextStore 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); diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java index 818930a5e3..1e63a346e5 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java @@ -1,6 +1,8 @@ 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. + * + *

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 + * + *

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 + * + *

As well as optional baggage items + * + *

On the other side of the communication a special Dispatcher is created when a message with + * DD_CONTEXT_CALL_ID is received. + * + *

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 knownConnections = + InstrumentationContext.get(Connection.class, Boolean.class); + + PROPAGATOR.attemptToPropagateContext(knownConnections, c, activeSpan); + } } } diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/StreamRemoteCallConstructorAdvice.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/StreamRemoteCallConstructorAdvice.java deleted file mode 100644 index dc457b979f..0000000000 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/client/StreamRemoteCallConstructorAdvice.java +++ /dev/null @@ -1,56 +0,0 @@ -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. - * - *

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 - * - *

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 - * - *

As well as optional baggage items - * - *

On the other side of the communication a special Dispatcher is created when a message with - * DD_CONTEXT_CALL_ID is received. - * - *

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 contextStore = - InstrumentationContext.get(Connection.class, Boolean.class); - - PROPAGATOR.attemptToPropagateContext(contextStore, c, activeSpan); - } -} diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java index d3a11bc177..c341331e5a 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ContextDispatcher.java @@ -1,24 +1,37 @@ 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 + * + *

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 + * + *

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 callableContextStore; - - public ContextDispatcher(final ContextStore 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 {} } diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ObjectTableAdvice.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ObjectTableAdvice.java deleted file mode 100644 index f187c2ef7f..0000000000 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/ObjectTableAdvice.java +++ /dev/null @@ -1,41 +0,0 @@ -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 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(); - } -} diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java index 3eb276bb7b..1ce41f9310 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/context/server/RmiServerContextInstrumentation.java @@ -1,6 +1,7 @@ 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 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(); + } } } diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerDecorator.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerDecorator.java index edd09cad90..155d64c7b2 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerDecorator.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerDecorator.java @@ -1,31 +1,11 @@ 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 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"}; diff --git a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerInstrumentation.java b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerInstrumentation.java index f479120483..5a96051c24 100644 --- a/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerInstrumentation.java +++ b/dd-java-agent/instrumentation/rmi/src/main/java/datadog/trace/instrumentation/rmi/server/RmiServerInstrumentation.java @@ -1,7 +1,9 @@ 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 contextStore() { - return Collections.singletonMap(Thread.class.getName(), AgentSpan.Context.class.getName()); - } - @Override public String[] helperClassNames() { return new String[] { @@ -55,22 +49,24 @@ public final class RmiServerInstrumentation extends Instrumenter.Default { @Override public Map, 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 threadContextStore = - InstrumentationContext.get(Thread.class, AgentSpan.Context.class); + final AgentSpan.Context context = THREAD_LOCAL_CONTEXT.getAndResetContext(); - final AgentSpan span = - DECORATE - .startSpanWithContext(threadContextStore) - .setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method)) - .setTag("span.origin.type", thiz.getClass().getCanonicalName()); + final AgentSpan span; + if (context == null) { + span = startSpan("rmi.request"); + } else { + span = startSpan("rmi.request", context); + } + + span.setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method)) + .setTag("span.origin.type", thiz.getClass().getCanonicalName()); DECORATE.afterStart(span); return activateSpan(span, true);