[rmi] use simple thread local for propagating context

+ add more documentation

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<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;
+  }
+}

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

@@ -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();
-  }
-}

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();
+    }
   }
 }

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<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();

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

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.
+ *
+ * <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);
+    }
   }
 }

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

@@ -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.
- *
- * <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,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
+ *
+ * <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 {}
 }

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

@@ -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<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();
-  }
-}

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<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();
+    }
   }
 }

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<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"};

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<String, String> 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<? 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 =
-          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);