diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java index 51650b3d2b2..ca75f489910 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java +++ b/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 sharedType = request == targetType ? null : types.find(name); + // existing type description from same classloader? - SharedTypeInfo 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) { diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTargetBytecodeTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTargetBytecodeTest.java new file mode 100644 index 00000000000..894ff5a5dd4 --- /dev/null +++ b/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(); + } +}