Test aot agent retransform

test/hotspot/jtreg/runtime/cds/appcds/aotCache/JavaAgent2.java

@@ -0,0 +1,184 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ *
+ */
+
+
+/**
+ * @test id=aot
+ * @summary -javaagent should be allowed in AOT workflow. However,
+ * classes transformed/redefined by agents will not be cached.
+ * @requires vm.cds.supports.aot.class.linking
+ * @library /test/lib /test/hotspot/jtreg/runtime/cds/appcds/test-classes
+ * @build JavaAgent2 JavaAgentRetransformer Util
+ * @run driver jdk.test.lib.helpers.ClassFileInstaller -jar app.jar JavaAgentApp2 JavaAgentApp2$ShouldNotBeTransformed JavaAgentApp2$ShouldBeTransformed
+ * @run driver JavaAgent2 AOT
+ */
+
+import java.util.random.RandomGenerator;
+import java.util.LinkedList;
+import java.util.Arrays;
+
+import jdk.test.lib.cds.CDSAppTester;
+import jdk.test.lib.process.OutputAnalyzer;
+import jdk.test.lib.helpers.ClassFileInstaller;
+
+public class JavaAgent2 {
+    static final String appJar = ClassFileInstaller.getJarPath("app.jar");
+    static final String mainClass = "JavaAgentApp2";
+
+    public static String agentClasses[] = {
+        "JavaAgentRetransformer",
+        "Util",
+    };
+    static String agentJar;
+
+    public static void main(String... args) throws Exception {
+        agentJar = ClassFileInstaller.writeJar("agent.jar",
+                                               ClassFileInstaller.Manifest.fromSourceFile("JavaAgentRetransformer.mf"),
+                                               agentClasses);
+
+        Tester t = new Tester();
+        t.run(args);
+    }
+
+    static class Tester extends CDSAppTester {
+        public Tester() {
+            super(mainClass);
+        }
+
+        @Override
+        public String classpath(RunMode runMode) {
+            return appJar;
+        }
+
+        @Override
+        public String[] vmArgs(RunMode runMode) {
+            switch (runMode) {
+            case RunMode.TRAINING:
+                return new String[] {
+                    "-javaagent:" + agentJar,
+                    "-Xlog:aot,cds",
+                    "-XX:+AOTClassLinking",
+                    // "-XX:CompileCommand=dontinline ShouldNotBeTransformed::doWork",
+                    "-XX:-TieredCompilation",
+                    "-XX:+PrintCompilation",
+                    "-XX:+UnlockDiagnosticVMOptions",
+                    "-XX:+PrintInlining",
+                };
+            case RunMode.ASSEMBLY:
+                return new String[] {
+                    "-javaagent:" + agentJar,
+                    "-Xlog:aot,cds",
+                    "-XX:+AOTClassLinking",
+                    "-XX:+AOTPrintTrainingInfo",
+                };
+            default:
+                return new String[] {
+                    "-javaagent:" + agentJar,
+                    "-Xlog:aot,cds",
+                    "-XX:+AOTClassLinking",
+                };
+            }
+        }
+
+        @Override
+        public String[] appCommandLine(RunMode runMode) {
+            return new String[] {
+                mainClass,
+            };
+        }
+
+        @Override
+        public void checkExecution(OutputAnalyzer out, RunMode runMode) throws Exception {
+            checkExecutionForAOTWorkflow(out, runMode);
+        }
+
+        static String agentLoadedMsg = "JavaAgentRetransformer.premain() is called";
+        static String agentPremainFinished = "JavaAgentRetransformer::premain() is finished";
+
+        public void checkExecutionForAOTWorkflow(OutputAnalyzer out, RunMode runMode) throws Exception {
+
+            if (runMode.isApplicationExecuted()) {
+                out.shouldContain(agentLoadedMsg);
+                out.shouldContain("Transforming: JavaAgentApp2$ShouldBeTransformed");
+                out.shouldContain(agentPremainFinished);
+            } else {
+                out.shouldNotContain(agentLoadedMsg);
+                out.shouldNotContain(agentPremainFinished);
+           }
+
+            switch (runMode) {
+            case RunMode.TRAINING:
+                out.shouldContain("Skipping JavaAgentApp2$ShouldBeTransformed: From ClassFileLoadHook");
+                out.shouldContain("Skipping JavaAgentApp2$ShouldBeTransformed: Has been redefined");
+                out.shouldContain("Skipping JavaAgentRetransformer: Unsupported location");
+                // should see compilation of $ShouldNotBeTransformed::doWork
+                out.shouldMatch("^[0-9]* *[0-9]* *[0-9] *JavaAgentApp2\\$ShouldNotBeTransformed::doWork");
+                // should see inlining of $ShouldBeTransformed::doWork
+                out.shouldMatch("JavaAgentApp2\\$ShouldBeTransformed::doWork \\([0-9]* bytes\\) *inline");
+                break;
+            case RunMode.ASSEMBLY:
+                out.shouldContain("Disabled all JVMTI agents during -XX:AOTMode=create");
+                // we should not be saving compiled training data for
+                // $ShouldBeTransformed.doWork as it depends on a
+                // method that was retransfromed
+                out.shouldNotMatch("^  C JavaAgentApp2\\$ShouldNotBeTransformed\\[.\\].doWork\\(\\)");
+                break;
+            }
+
+        }
+    }
+}
+
+class JavaAgentApp2 {
+    public static void main(String[] args) {
+        ShouldNotBeTransformed s = new ShouldNotBeTransformed();
+        // Force profile/compile of ShouldNotBeTransformed::dowork.
+        // The compiler should inline ShouldBeTransformed::doWork
+        for (int i = 0; i < 10_000; i++) {
+            s.doWork();
+        }
+        // transform ShouldBeTransformed::doWork to a simple return
+        JavaAgentRetransformer.doRetransform(ShouldBeTransformed.class);
+        // Profile and compile again
+        for (int i = 0; i < 10_000; i++) {
+            s.doWork();
+        }
+    }
+
+    static class ShouldNotBeTransformed {
+        private ShouldBeTransformed r = new ShouldBeTransformed();
+        private int value = 0;
+        public void doWork() {
+            value = r.doWork(value);
+        }
+        public int getValue() { return value; }
+    }
+
+    static class ShouldBeTransformed {
+        private RandomGenerator random = RandomGenerator.of("L64X128MixRandom");
+        public int doWork(int value) {
+            return (value * (random.nextInt(5) + 1) + random.nextInt(7));
+        }
+    }
+}

test/hotspot/jtreg/runtime/cds/appcds/aotCache/JavaAgentRetransformer.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ *
+ */
+
+import java.lang.classfile.ClassFile;
+import java.lang.classfile.MethodModel;
+import java.lang.classfile.ClassTransform;
+import java.lang.constant.ConstantDescs;
+import java.lang.constant.ClassDesc;
+import java.lang.constant.MethodTypeDesc;
+
+import java.lang.System.Logger.Level;
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.IllegalClassFormatException;
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.AccessFlag;
+import java.security.ProtectionDomain;
+
+// This class is available on the classpath so it can be accessed by JavaAgentRetransformApp
+public class JavaAgentRetransformer  implements ClassFileTransformer {
+    private static Instrumentation savedInstrumentation;
+    private static final System.Logger LOGGER = System.getLogger(JavaAgentRetransformer.class.getName());
+    private static Class<?> current = null;
+    public static void premain(String agentArguments, Instrumentation instrumentation) {
+        System.out.println("JavaAgentRetransformer.premain() is called");
+        instrumentation.addTransformer(new JavaAgentRetransformer(), /*canRetransform=*/true);
+        savedInstrumentation = instrumentation;
+
+        LOGGER.log(Level.WARNING, "JavaAgentRetransformer::premain() is finished");
+    }
+
+    public static Instrumentation getInstrumentation() {
+        return savedInstrumentation;
+    }
+
+    public static void agentmain(String args, Instrumentation inst) throws Exception {
+        premain(args, inst);
+    }
+
+    public static void doRetransform(Class<?> clazz) {
+        current = clazz;
+        Class<?> classes[] = new Class<?>[1];
+        classes[0] = clazz;
+        try {
+            savedInstrumentation.retransformClasses(classes);
+        } catch (Throwable t) {
+            t.printStackTrace();
+        }
+    }
+
+    public byte[] transform(ClassLoader loader, String name, Class<?> classBeingRedefined,
+                            ProtectionDomain pd, byte[] buffer) throws IllegalClassFormatException {
+        // only accept retransform requests for the current class
+        if (classBeingRedefined == null) {
+            return null;
+        }
+        if (classBeingRedefined == current) {
+            System.out.println("Transforming: " + name);
+            try {
+                buffer = transformDoWork(buffer);
+            } catch (Throwable t) {
+                t.printStackTrace();
+            }
+            Thread.dumpStack();
+            return buffer;
+        }
+        return null;
+    }
+
+    final static ClassDesc CD_int = ClassDesc.ofDescriptor("I");
+    final static MethodTypeDesc MTD_int_int = MethodTypeDesc.of(CD_int, CD_int);
+
+    // replace the body of any public void instance method that looks
+    // like "int doWork(int arg0)" with "return arg0"
+    static byte[] transformDoWork(byte[] buffer) {
+        ClassTransform t = (builder, element) -> {
+            if (element instanceof MethodModel method &&
+                method.methodName().equalsString("doWork") &&
+                method.methodTypeSymbol() == MTD_int_int &&
+                method.flags().has(AccessFlag.PUBLIC) &&
+                !method.flags().has(AccessFlag.STATIC)) {
+                builder.withMethodBody("doWork",
+                                       MTD_int_int,
+                                       ClassFile.ACC_PUBLIC,
+                                       cob -> cob.iload(0).ireturn());
+            } else {
+                builder.with(element);  // leaves the element in place
+            }
+        };
+        var c = ClassFile.of();
+        byte[] newBytes = c.transformClass(c.parse(buffer), t);
+        return newBytes;
+    }
+}

test/hotspot/jtreg/runtime/cds/appcds/aotCache/JavaAgentRetransformer.mf

@@ -0,0 +1,5 @@
+Manifest-Version: 1.0
+Premain-Class: JavaAgentRetransformer
+Agent-Class: JavaAgentRetransformer
+Can-Retransform-Classes: true
+Can-Redefine-Classes: true

test/hotspot/jtreg/runtime/cds/appcds/aotCache/TEST.properties

@@ -0,0 +1 @@
+maxOutputSize = 1000000