Applying suggestions and use ClassValue

Author: amarziali

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

@@ -139,8 +139,12 @@ 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
-  testImplementation 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')

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

@@ -5,10 +5,8 @@
 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;
-import net.bytebuddy.asm.Advice;
-import net.bytebuddy.implementation.bytecode.assign.Assigner;
-import org.mule.runtime.tracer.api.EventTracer;
 
 @AutoService(InstrumenterModule.class)
 public class JpmsMuleInstrumentation extends InstrumenterModule.Tracing
@@ -39,20 +37,23 @@ public String[] helperClassNames() {
   }
 
   @Override
-  public void methodAdvice(MethodTransformer transformer) {
-    // it does not work with typeInitializer()
-    transformer.applyAdvice(isConstructor(), getClass().getName() + "$JpmsClearanceAdvice");
+  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(),
+    };
   }
 
-  public static class JpmsClearanceAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void openOnReturn(@Advice.This(typing = Assigner.Typing.DYNAMIC) Object self) {
-      JpmsAdvisingHelper.allowAccessOnModuleClass(self.getClass());
-    }
-
-    private static void muzzleCheck(final EventTracer<?> tracer) {
-      // introduced in 4.5.0
-      tracer.endCurrentSpan(null);
-    }
+  @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

@@ -1,22 +1,11 @@
 package datadog.trace.instrumentation.mule4;
 
-import java.util.WeakHashMap;
+import datadog.trace.api.GenericClassValue;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class JpmsAdvisingHelper {
-  private static final WeakHashMap<Class<?>, Boolean> ALREADY_PROCESSED_CACHE = new WeakHashMap<>();
-
-  public static void allowAccessOnModuleClass(final Class<?> cls) {
-    if (Boolean.TRUE.equals(ALREADY_PROCESSED_CACHE.putIfAbsent(cls, Boolean.TRUE))) {
-      return;
-    }
-    final Module module = cls.getModule();
-    if (module != null) {
-      try {
-        module.addExports(cls.getPackageName(), module.getClassLoader().getUnnamedModule());
-      } catch (Throwable ignored) {
-      }
-    }
-  }
+  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) {
+        }
+      }
+    }
+  }
+}