Mule: do not crash with JPMS

Author: amarziali

dd-java-agent/instrumentation/mule-4/build.gradle

@@ -43,6 +43,7 @@ muzzle {
 }
 
 apply from: "$rootDir/gradle/java.gradle"
+apply plugin: "idea"
 
 addTestSuiteForDir('mule46ForkedTest', 'test')
 addTestSuiteForDir('latestDepForkedTest', 'test')
@@ -81,6 +82,9 @@ configurations.all {
 }
 
 sourceSets {
+  main_java11 {
+    java.srcDirs "${project.projectDir}/src/main/java11"
+  }
   test {
     output.dir("$buildDir/generated-resources/test", builtBy: 'generateAppResources')
   }
@@ -92,6 +96,20 @@ sourceSets {
   }
 }
 
+compileMain_java11Java.configure {
+  setJavaVersion(it, 11)
+  sourceCompatibility = JavaVersion.VERSION_1_8
+  targetCompatibility = JavaVersion.VERSION_1_8
+}
+
+jar {
+  from sourceSets.main_java11.output
+}
+
+forbiddenApisMain_java11 {
+  failOnMissingClasses = false
+}
+
 tasks.named("compileTestGroovy").configure {
   dependsOn 'mvnPackage', 'extractMuleServices'
 }
@@ -112,7 +130,10 @@ tasks.named("compileLatestDepForkedTestJava").configure {
 dependencies {
   compileOnly group: 'org.mule.runtime', name: 'mule-core', version: muleVersion
   compileOnly group: 'org.mule.runtime', name: 'mule-tracer-customization-impl', version: muleVersion
-
+  main_java11CompileOnly project(':internal-api')
+  main_java11CompileOnly project(':dd-java-agent:agent-tooling')
+  main_java11CompileOnly project(':dd-java-agent:agent-bootstrap')
+  main_java11CompileOnly sourceSets.main.output
   testImplementation project(':dd-java-agent:instrumentation:aws-common')
   testImplementation project(':dd-java-agent:instrumentation:reactor-core-3.1')
   testImplementation project(':dd-java-agent:instrumentation:reactive-streams')
@@ -234,3 +255,9 @@ spotless {
     target "**/*.java"
   }
 }
+
+idea {
+  module {
+    jdkName = '11'
+  }
+}

dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/JpmsMuleInstrumentation.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.mule4;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Platform;
+
+@AutoService(InstrumenterModule.class)
+public class JpmsMuleInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.HasMethodAdvice, Instrumenter.ForKnownTypes {
+  public JpmsMuleInstrumentation() {
+    super("mule", "mule-jpms");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && Platform.isJavaVersionAtLeast(9);
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      // same module but they can be initialized in any order
+      "org.mule.runtime.tracer.customization.impl.info.ExecutionInitialSpanInfo",
+      "org.mule.runtime.tracer.customization.impl.provider.LazyInitialSpanInfo",
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".JpmsAdvisingHelper",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // it does not work with typeInitializer()
+    transformer.applyAdvice(isConstructor(), packageName + ".JpmsClearanceAdvice");
+  }
+}

dd-java-agent/instrumentation/mule-4/src/main/java11/datadog/trace/instrumentation/mule4/JpmsAdvisingHelper.java

@@ -0,0 +1,13 @@
+package datadog.trace.instrumentation.mule4;
+
+import java.util.WeakHashMap;
+
+public class JpmsAdvisingHelper {
+  private static final WeakHashMap<Module, Boolean> ALREADY_PROCESSED_CACHE = new WeakHashMap<>();
+
+  public static boolean isModuleAlreadyProcessed(final Module module) {
+    return Boolean.TRUE.equals(ALREADY_PROCESSED_CACHE.putIfAbsent(module, Boolean.TRUE));
+  }
+
+  private JpmsAdvisingHelper() {}
+}

dd-java-agent/instrumentation/mule-4/src/main/java11/datadog/trace/instrumentation/mule4/JpmsClearanceAdvice.java

@@ -0,0 +1,20 @@
+package datadog.trace.instrumentation.mule4;
+
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+
+public class JpmsClearanceAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void openOnReturn(@Advice.This(typing = Assigner.Typing.DYNAMIC) Object self) {
+    final Module module = self.getClass().getModule();
+    if (module == null || JpmsAdvisingHelper.isModuleAlreadyProcessed(module)) {
+      return;
+    }
+    for (String pn : module.getPackages()) {
+      try {
+        module.addExports(pn, module.getClassLoader().getUnnamedModule());
+      } catch (Throwable ignored) {
+      }
+    }
+  }
+}