Merge pull request #3 from laurit/rmi-open

share module opening between indy and inline

Author: SylvainJuge

instrumentation/rmi/javaagent/build.gradle.kts

@@ -34,14 +34,5 @@ tasks {
   }
   withType<Test>().configureEach {
     jvmArgs("-Djava.rmi.server.hostname=127.0.0.1")
-
-    // Can only export on Java 9+
-    val testJavaVersion =
-      gradle.startParameter.projectProperties.get("testJavaVersion")?.let(JavaVersion::toVersion)
-        ?: JavaVersion.current()
-    if (testJavaVersion.isJava9Compatible) {
-      jvmArgs("--add-exports=java.rmi/sun.rmi.server=ALL-UNNAMED")
-      jvmArgs("--add-exports=java.rmi/sun.rmi.transport=ALL-UNNAMED")
-    }
   }
 }

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/RmiContextPropagationInstrumentationModule.java

@@ -13,10 +13,12 @@
 import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.instrumentation.rmi.context.client.RmiClientContextInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.rmi.context.server.RmiServerContextInstrumentation;
+import java.rmi.Remote;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import net.bytebuddy.utility.JavaModule;
 
 @AutoService(InstrumentationModule.class)
 public class RmiContextPropagationInstrumentationModule extends InstrumentationModule
@@ -31,9 +33,8 @@ public List<TypeInstrumentation> typeInstrumentations() {
   }
 
   @Override
-  public Map<String, List<String>> jpmsModulesToOpen() {
-    String witnessClass = "sun.rmi.transport.StreamRemoteCall";
+  public Map<JavaModule, List<String>> jpmsModulesToOpen() {
     return Collections.singletonMap(
-        witnessClass, Arrays.asList("sun.rmi.server", "sun.rmi.transport"));
+        JavaModule.ofType(Remote.class), Arrays.asList("sun.rmi.server", "sun.rmi.transport"));
   }
 }

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/client/RmiClientContextInstrumentation.java

@@ -13,19 +13,13 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
-import io.opentelemetry.javaagent.bootstrap.AgentClassLoader;
-import io.opentelemetry.javaagent.bootstrap.InstrumentationHolder;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.lang.instrument.Instrumentation;
 import java.rmi.server.ObjID;
-import java.util.Collections;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.dynamic.loading.ClassInjector;
 import net.bytebuddy.matcher.ElementMatcher;
-import net.bytebuddy.utility.JavaModule;
 import sun.rmi.transport.Connection;
 
 /**
@@ -63,30 +57,6 @@ public void transform(TypeTransformer transformer) {
             .and(takesArgument(0, named("sun.rmi.transport.Connection")))
             .and(takesArgument(1, named("java.rmi.server.ObjID"))),
         getClass().getName() + "$StreamRemoteCallConstructorAdvice");
-
-    // expose sun.rmi.transport.StreamRemoteCall to helper classes
-    transformer.applyTransformer(
-        (builder, typeDescription, classLoader, javaModule, protectionDomain) -> {
-          if (JavaModule.isSupported()
-              && classLoader == null
-              && "sun.rmi.transport.StreamRemoteCall".equals(typeDescription.getName())
-              && javaModule != null) {
-            Instrumentation instrumentation = InstrumentationHolder.getInstrumentation();
-            ClassInjector.UsingInstrumentation.redefineModule(
-                instrumentation,
-                javaModule,
-                Collections.emptySet(),
-                Collections.emptyMap(),
-                Collections.singletonMap(
-                    "sun.rmi.transport",
-                    // AgentClassLoader is in unnamed module of the bootstrap class loader which is
-                    // where helper classes are also
-                    Collections.singleton(JavaModule.ofType(AgentClassLoader.class))),
-                Collections.emptySet(),
-                Collections.emptyMap());
-          }
-          return builder;
-        });
   }
 
   @SuppressWarnings("unused")

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/jpms/ExposeRmiModuleInstrumentation.java

@@ -12,6 +12,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * This class is internal and is hence not for public use. Its APIs are unstable and can change at
@@ -70,9 +71,9 @@ default List<String> agentPackagesToHide() {
    * to get a reference to the module. <br>
    * Map value is a list of packages to open in the module
    *
-   * @return map of "witness class" FQN as key, list of packages as value.
+   * @return map of module to open as key, list of packages as value.
    */
-  default Map<String, List<String>> jpmsModulesToOpen() {
+  default Map<JavaModule, List<String>> jpmsModulesToOpen() {
     return Collections.emptyMap();
   }
 }

instrumentation/rmi/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rmi/context/jpms/RmiJpmsInstrumentationModule.java

@@ -8,14 +8,17 @@
 import static java.util.logging.Level.FINE;
 import static java.util.logging.Level.WARNING;
 
+import io.opentelemetry.javaagent.bootstrap.AgentClassLoader;
 import java.lang.instrument.Instrumentation;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.logging.Logger;
+import net.bytebuddy.description.type.PackageDescription;
+import net.bytebuddy.dynamic.loading.ClassInjector;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * Module opener provides ability to open JPMS modules and allows instrumentation classloader to
@@ -28,27 +31,30 @@ public class ModuleOpener {
 
   private static final Logger logger = Logger.getLogger(ModuleOpener.class.getName());
 
+  // AgentClassLoader is in unnamed module of the bootstrap loader
+  private static final JavaModule UNNAMED_BOOT_MODULE = JavaModule.ofType(AgentClassLoader.class);
+
   private ModuleOpener() {}
 
   /**
    * Opens JPMS module to a class loader unnamed module
    *
-   * @param classFromTargetModule class from target module
-   * @param openTo class loader to open module for
+   * @param targetModule target module
+   * @param openTo class loader to open module for, {@literal null} to use the unnamed module of bootstrap classloader.
    * @param packagesToOpen packages to open
    */
   public static void open(
       Instrumentation instrumentation,
-      Class<?> classFromTargetModule,
-      ClassLoader openTo,
+      JavaModule targetModule,
+      @Nullable ClassLoader openTo,
       Collection<String> packagesToOpen) {
 
-    Module targetModule = classFromTargetModule.getModule();
-    Module openToModule = openTo.getUnnamedModule();
-    Set<Module> openToModuleSet = Collections.singleton(openToModule);
-    Map<String, Set<Module>> missingOpens = new HashMap<>();
+    JavaModule openToModule =
+        openTo != null ? JavaModule.of(openTo.getUnnamedModule()) : UNNAMED_BOOT_MODULE;
+    Set<JavaModule> openToModuleSet = Collections.singleton(openToModule);
+    Map<String, Set<JavaModule>> missingOpens = new HashMap<>();
     for (String packageName : packagesToOpen) {
-      if (!targetModule.isOpen(packageName, openToModule)) {
+      if (!targetModule.isOpened(new PackageDescription.Simple(packageName), openToModule)) {
         missingOpens.put(packageName, openToModuleSet);
         logger.log(
             FINE,
@@ -60,22 +66,17 @@ public static void open(
       return;
     }
 
-    if (!instrumentation.isModifiableModule(targetModule)) {
-      logger.log(WARNING, "Module '{0}' can't be modified", targetModule);
-      return;
-    }
-
     try {
-      instrumentation.redefineModule(
+      ClassInjector.UsingInstrumentation.redefineModule(
+          instrumentation,
           targetModule,
-          Collections.<Module>emptySet(), // reads
-          Collections.<String, Set<Module>>emptyMap(), // exports
-          missingOpens, // opens
-          Collections.<Class<?>>emptySet(), // uses
-          Collections.<Class<?>, List<Class<?>>>emptyMap() // provides
-          );
+          Collections.emptySet(),
+          Collections.emptyMap(),
+          missingOpens,
+          Collections.emptySet(),
+          Collections.emptyMap());
     } catch (Exception e) {
-      logger.log(WARNING, "unable to redefine module", e);
+      logger.log(WARNING, "Failed to redefine module '" + targetModule.getActualName() + "'", e);
     }
   }
 }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java

@@ -17,6 +17,7 @@
 import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
 import io.opentelemetry.javaagent.tooling.HelperClassDefinition;
 import io.opentelemetry.javaagent.tooling.HelperInjector;
+import io.opentelemetry.javaagent.tooling.ModuleOpener;
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.bytebuddy.LoggingFailSafeMatcher;
@@ -36,6 +37,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationSource;
@@ -202,13 +204,31 @@ private AgentBuilder installInjectingModule(
     VirtualFieldImplementationInstaller contextProvider =
         virtualFieldInstallerFactory.create(instrumentationModule);
 
+    AtomicBoolean openerRun = new AtomicBoolean();
     AgentBuilder agentBuilder = parentAgentBuilder;
     for (TypeInstrumentation typeInstrumentation : typeInstrumentations) {
-
       AgentBuilder.Identified.Extendable extendableAgentBuilder =
           setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
               .and(muzzleMatcher)
               .transform(ConstantAdjuster.instance())
+              .transform(
+                  (builder, typeDescription, classLoader, module, protectionDomain) -> {
+                    if (instrumentationModule instanceof ExperimentalInstrumentationModule
+                        && !openerRun.get()) {
+                      ExperimentalInstrumentationModule experimentalModule =
+                          (ExperimentalInstrumentationModule) instrumentationModule;
+                      experimentalModule
+                          .jpmsModulesToOpen()
+                          .forEach(
+                              (javaModule, packages) -> {
+                                ModuleOpener.open(
+                                    instrumentation, javaModule, classLoader, packages);
+                              });
+                      openerRun.set(true);
+                    }
+
+                    return builder;
+                  })
               .transform(helperInjector);
       extendableAgentBuilder = contextProvider.injectHelperClasses(extendableAgentBuilder);
       extendableAgentBuilder = contextProvider.rewriteVirtualFieldsCalls(extendableAgentBuilder);

javaagent-tooling/javaagent-tooling-java9/src/main/java/io/opentelemetry/javaagent/tooling/ModuleOpener.java

@@ -94,16 +94,8 @@ public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
         experimentalModule
             .jpmsModulesToOpen()
             .forEach(
-                (className, packages) -> {
-                  Class<?> type;
-                  try {
-                    type = Class.forName(className, false, instrumentedClassLoader);
-                  } catch (ClassNotFoundException e) {
-                    throw new IllegalStateException("missing witness class " + className, e);
-                  }
-
-                  ModuleOpener.open(instrumentation, type, loader, packages);
-                });
+                (javaModule, packages) ->
+                    ModuleOpener.open(instrumentation, javaModule, loader, packages));
       }
     }
     return loader;