Mule: do not crash with JPMS (#8187)

* Mule: do not crash with JPMS

* muzzle

* Applying suggestions and use ClassValue

Author: amarziali

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

@@ -35,14 +35,23 @@ muzzle {
     module = 'mule-core'
     versions = "[$muleVersion,)"
     javaVersion = "17"
-    assertInverse true
     excludeDependency 'om.google.guava:guava'
     excludeDependency 'com.google.code.findbugs:jsr305'
     additionalDependencies +="org.mule.runtime:mule-tracer-customization-impl:$muleVersion"
   }
+
+  fail {
+    name = 'before-4.5.0'
+    group = 'org.mule.runtime'
+    module = 'mule-core'
+    versions =  "[, $muleVersion)"
+    excludeDependency 'om.google.guava:guava'
+    excludeDependency 'com.google.code.findbugs:jsr305'
+  }
 }
 
 apply from: "$rootDir/gradle/java.gradle"
+apply plugin: "idea"
 
 addTestSuiteForDir('mule46ForkedTest', 'test')
 addTestSuiteForDir('latestDepForkedTest', 'test')
@@ -81,6 +90,9 @@ configurations.all {
 }
 
 sourceSets {
+  main_java11 {
+    java.srcDirs "${project.projectDir}/src/main/java11"
+  }
   test {
     output.dir("$buildDir/generated-resources/test", builtBy: 'generateAppResources')
   }
@@ -92,6 +104,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 +138,13 @@ 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
+  compileOnly sourceSets.main_java11.output
 
+  main_java11CompileOnly project(':internal-api')
+  main_java11CompileOnly project(':dd-java-agent:agent-tooling')
+  main_java11CompileOnly project(':dd-java-agent:agent-bootstrap')
+
+  testImplementation sourceSets.main_java11.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 +266,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,59 @@
+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.agent.tooling.muzzle.Reference;
+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 Reference[] additionalMuzzleReferences() {
+    return new Reference[] {
+      // added in 4.5.0
+      new Reference.Builder("org.mule.runtime.tracer.api.EventTracer")
+          .withMethod(
+              new String[0],
+              Reference.EXPECTS_NON_STATIC | Reference.EXPECTS_PUBLIC,
+              "endCurrentSpan",
+              "V",
+              "Lorg/mule/runtime/api/event/Event;")
+          .build(),
+    };
+  }
+
+  @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,11 @@
+package datadog.trace.instrumentation.mule4;
+
+import datadog.trace.api.GenericClassValue;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class JpmsAdvisingHelper {
+  public static final ClassValue<AtomicBoolean> ALREADY_PROCESSED_CACHE =
+      GenericClassValue.constructing(AtomicBoolean.class);
+
+  private JpmsAdvisingHelper() {}
+}

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

@@ -0,0 +1,26 @@
+package datadog.trace.instrumentation.mule4;
+
+import static datadog.trace.instrumentation.mule4.JpmsAdvisingHelper.ALREADY_PROCESSED_CACHE;
+
+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 Class<?> cls = self.getClass();
+    if (ALREADY_PROCESSED_CACHE.get(cls).compareAndSet(false, true)) {
+      final Module module = cls.getModule();
+      if (module != null) {
+        try {
+          // This call needs imperatively to be done from the same module we're adding exports
+          // because the jdk is checking that the caller belongs to the same module.
+          // The code of this advice is getting inlined into the constructor of the class belonging
+          // to that package so it will work. Moving the same to a helper won't.
+          module.addExports(cls.getPackageName(), module.getClassLoader().getUnnamedModule());
+        } catch (Throwable ignored) {
+        }
+      }
+    }
+  }
+}