🍒 13109 - Fix stale type descriptions when transforming multi-release classes

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java

@@ -108,9 +108,7 @@ final class TypeFactory {
 
   ClassFileLocator classFileLocator;
 
-  String targetName;
-
-  byte[] targetBytecode;
+  private LazyType targetType;
 
   /** Sets the current class-loader context of this type-factory. */
   void switchContext(ClassLoader classLoader) {
@@ -151,8 +149,8 @@ static void clear() {
    * used for matching and restore it afterwards.
    */
   void beginTransform(String name, byte[] bytecode) {
-    targetName = name;
-    targetBytecode = bytecode;
+    targetType = new LazyType(name);
+    targetType.location = new ClassFileLocator.Resolution.Explicit(bytecode);
 
     if (installing) {
       originalClassLoader = currentClassLoader;
@@ -173,7 +171,7 @@ boolean disableFullDescriptions() {
 
   /** Cleans-up local caches to minimise memory use once we're done with the type-factory. */
   void endTransform() {
-    if (null == targetName) {
+    if (null == targetType) {
       return; // transformation didn't reach resolve step
     }
 
@@ -184,8 +182,7 @@ void endTransform() {
       clearReferences();
     }
 
-    targetName = null;
-    targetBytecode = null;
+    targetType = null;
     createOutlines = OUTLINING_ENABLED;
   }
 
@@ -228,12 +225,19 @@ static TypeDescription findType(String name) {
   }
 
   private TypeDescription deferTypeResolution(String name) {
+    if (null != targetType && name.equals(targetType.name)) {
+      return targetType; // keep target type alive during request
+    }
     return deferredTypes.computeIfAbsent(name, deferType);
   }
 
   /** Attempts to resolve the named type using the current context. */
   TypeDescription resolveType(LazyType request) {
     if (null != classFileLocator) {
+      // deferred references issued before beginTransform should reuse the target's description
+      if (null != targetType && request != targetType && request.name.equals(targetType.name)) {
+        return targetType.doResolve(false);
+      }
       TypeDescription result;
       if (createOutlines) {
         if ("java.lang.Object".equals(request.name)) {
@@ -258,8 +262,10 @@ private TypeDescription lookupType(
     boolean isOutline = typeParser == outlineTypeParser;
     long fromTick = InstrumenterMetrics.tick();
 
+    // the JVM's bytes are authoritative for the target; its LazyType reuses the parsed result
+    SharedTypeInfo<TypeDescription> sharedType = request == targetType ? null : types.find(name);
+
     // existing type description from same classloader?
-    SharedTypeInfo<TypeDescription> sharedType = types.find(name);
     if (null != sharedType
         && (name.startsWith("java.") || sharedType.sameClassLoader(classLoaderId))) {
       InstrumenterMetrics.reuseTypeDescription(fromTick, isOutline);
@@ -369,9 +375,6 @@ public byte[] getBytecode() {
     }
 
     private ClassFileLocator.Resolution locateClassFile() {
-      if (name.equals(targetName)) {
-        return new ClassFileLocator.Resolution.Explicit(targetBytecode);
-      }
       try {
         return classFileLocator.locate(name);
       } catch (Throwable ignored) {

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTargetBytecodeTest.java

@@ -0,0 +1,108 @@
+package datadog.trace.agent.tooling.bytebuddy.outline;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+import java.io.IOException;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import net.bytebuddy.description.type.TypeDescription;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.Opcodes;
+
+class TypeFactoryTargetBytecodeTest {
+  private static final String TYPE_NAME = "example.MultiReleaseType";
+
+  @TempDir Path resources;
+
+  @AfterEach
+  void clearTypes() {
+    TypeFactory.typeFactory.get().endTransform();
+    TypeFactory.typeFactory.get().endInstall();
+    TypeFactory.typeFactory.remove();
+    TypeFactory.clear();
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {false, true})
+  void transformBytesOverridePreviouslyResolvedResource(boolean fullDescriptions) throws Exception {
+    try (URLClassLoader loader = resourceLoader()) {
+      TypeFactory factory = TypeFactory.typeFactory.get();
+      factory.beginInstall();
+      factory.switchContext(loader);
+      factory.beginTransform("example.OtherType", new byte[0]);
+      assertEquals("java.util.ArrayList", resolveParent(factory, fullDescriptions));
+      factory.endTransform();
+
+      // Model a loader whose resource returns the base entry but defines the versioned entry.
+      factory.switchContext(loader);
+      factory.beginTransform(TYPE_NAME, classBytes("java/util/LinkedList"));
+      assertEquals("java.util.LinkedList", resolveParent(factory, fullDescriptions));
+
+      TypeFactory.LazyType target = (TypeFactory.LazyType) TypeFactory.findType(TYPE_NAME);
+      TypeDescription description = target.doResolve(true);
+      // Evict the target from the small general-purpose cache without resolving these names.
+      for (int i = 0; i < 100; i++) {
+        TypeFactory.findType("example.OtherType" + i);
+      }
+      TypeFactory.LazyType repeated = (TypeFactory.LazyType) TypeFactory.findType(TYPE_NAME);
+      assertSame(description, repeated.doResolve(true));
+
+      factory.endTransform();
+      factory.switchContext(loader);
+      factory.beginTransform(TYPE_NAME, classBytes("java/util/HashSet"));
+      assertEquals("java.util.HashSet", resolveParent(factory, fullDescriptions));
+    }
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {false, true})
+  void previouslyIssuedReferenceReusesTargetDescription(boolean fullDescriptions) throws Exception {
+    try (URLClassLoader loader = resourceLoader()) {
+      TypeFactory factory = TypeFactory.typeFactory.get();
+      factory.beginInstall();
+      factory.switchContext(loader);
+      TypeFactory.LazyType retained = (TypeFactory.LazyType) TypeFactory.findType(TYPE_NAME);
+
+      factory.beginTransform(TYPE_NAME, classBytes("java/util/LinkedList"));
+      if (fullDescriptions) {
+        factory.enableFullDescriptions();
+      }
+      TypeDescription description = retained.doResolve(true);
+      assertEquals("java.util.LinkedList", description.getSuperClass().asErasure().getName());
+      TypeFactory.LazyType target = (TypeFactory.LazyType) TypeFactory.findType(TYPE_NAME);
+      assertSame(description, target.doResolve(true));
+    }
+  }
+
+  private URLClassLoader resourceLoader() throws IOException {
+    Path resource = resources.resolve(TYPE_NAME.replace('.', '/') + ".class");
+    Files.createDirectories(resource.getParent());
+    Files.write(resource, classBytes("java/util/ArrayList"));
+    return new URLClassLoader(new URL[] {resources.toUri().toURL()}, getClass().getClassLoader());
+  }
+
+  private static String resolveParent(TypeFactory factory, boolean fullDescriptions) {
+    if (fullDescriptions) {
+      factory.enableFullDescriptions();
+    }
+    TypeDescription type = TypeFactory.findType(TYPE_NAME);
+    if (fullDescriptions) {
+      type.getDeclaredMethods();
+    }
+    return type.getSuperClass().asErasure().getName();
+  }
+
+  private static byte[] classBytes(String parent) {
+    ClassWriter writer = new ClassWriter(0);
+    writer.visit(Opcodes.V1_8, Opcodes.ACC_PUBLIC, TYPE_NAME.replace('.', '/'), null, parent, null);
+    writer.visitEnd();
+    return writer.toByteArray();
+  }
+}