Improve compatibility with other agents (#7916)

Fixes the issue described in #7887 Hopefully resolves #7594 We should not use cached `TypeDescription` for currently transformed class as the cached description will not be the same as newly created one if the class bytes were transformed. For example if another agent adds an interface to the class then returning the cached description that does not have that interface would result in bytebuddy removing that interface, see https://github.com/raphw/byte-buddy/blob/665a090c733c37339788cc5a1c441bd47fb77ef0/byte-buddy-dep/src/main/java/net/bytebuddy/dynamic/scaffold/TypeWriter.java#L5012
open-telemetry · Mar 1, 2023 · dd32ff3 · dd32ff3
1 parent 6f03734
commit dd32ff3
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 0 deletions.
diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
@@ -138,6 +138,7 @@ private static void installBytebuddyAgent(
             .with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
             .with(AgentTooling.poolStrategy())
             .with(new ClassLoadListener())
+            .with(AgentTooling.transformListener())
             .with(AgentTooling.locationStrategy());
     if (JavaModule.isSupported()) {
       agentBuilder = agentBuilder.with(new ExposeAgentBootstrapListener(inst));

diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategy.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategy.java
@@ -275,6 +275,15 @@ public TypePool.Resolution find(String className) {
       if (OBJECT_NAME.equals(className)) {
         return OBJECT_RESOLUTION;
       }
+      // Skip cache for the type that is currently being transformed.
+      // If class has been transformed by another agent or by class loader it is possible that the
+      // cached TypeDescription isn't the same as the one built from the actual bytes that are
+      // being defined. For example if another agent adds an interface to the class then returning
+      // the cached description that does not have that interface would result in bytebuddy removing
+      // that interface.
+      if (AgentTooling.isTransforming(loaderRef != null ? loaderRef.get() : null, className)) {
+        return null;
+      }
 
       TypePool.Resolution existingResolution =
           sharedResolutionCache.get(new TypeCacheKey(loaderHash, loaderRef, className));

diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentTooling.java
@@ -8,6 +8,7 @@
 import java.util.Iterator;
 import java.util.ServiceLoader;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * This class contains class references for objects shared by the agent installer as well as muzzle
@@ -22,6 +23,8 @@ public final class AgentTooling {
   private static final AgentBuilder.PoolStrategy POOL_STRATEGY =
       new AgentCachingPoolStrategy(LOCATION_STRATEGY);
 
+  private static final ThreadLocal<CurrentTransform> CURRENT_TRANSFORM = new ThreadLocal<>();
+
   public static AgentLocationStrategy locationStrategy() {
     return LOCATION_STRATEGY;
   }
@@ -30,6 +33,10 @@ public static AgentBuilder.PoolStrategy poolStrategy() {
     return POOL_STRATEGY;
   }
 
+  public static AgentBuilder.Listener transformListener() {
+    return new ClassTransformListener();
+  }
+
   private static ClassLoader getBootstrapProxy() {
     Iterator<BootstrapProxyProvider> iterator =
         ServiceLoader.load(BootstrapProxyProvider.class, AgentTooling.class.getClassLoader())
@@ -42,5 +49,38 @@ private static ClassLoader getBootstrapProxy() {
     return null;
   }
 
+  public static boolean isTransforming(ClassLoader classLoader, String className) {
+    CurrentTransform currentTransform = CURRENT_TRANSFORM.get();
+    if (currentTransform == null) {
+      return false;
+    }
+    return currentTransform.className.equals(className)
+        && currentTransform.classLoader == classLoader;
+  }
+
+  private static class ClassTransformListener extends AgentBuilder.Listener.Adapter {
+    @Override
+    public void onDiscovery(
+        String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
+      CURRENT_TRANSFORM.set(new CurrentTransform(classLoader, typeName));
+    }
+
+    @Override
+    public void onComplete(
+        String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
+      CURRENT_TRANSFORM.remove();
+    }
+  }
+
+  private static class CurrentTransform {
+    private final ClassLoader classLoader;
+    private final String className;
+
+    CurrentTransform(ClassLoader classLoader, String className) {
+      this.classLoader = classLoader;
+      this.className = className;
+    }
+  }
+
   private AgentTooling() {}
 }