Merge pull request #1543 from codeflash-ai/fix/java/line-profiler

feat: implement line-level profiling agent with ASM instrumentation

Author: HeshamHM28

codeflash-java-runtime/pom.xml

@@ -48,6 +48,18 @@
             <version>3.45.0.0</version>
         </dependency>
 
+        <!-- ASM for bytecode instrumentation (profiler agent) -->
+        <dependency>
+            <groupId>org.ow2.asm</groupId>
+            <artifactId>asm</artifactId>
+            <version>9.7.1</version>
+        </dependency>
+        <dependency>
+            <groupId>org.ow2.asm</groupId>
+            <artifactId>asm-commons</artifactId>
+            <version>9.7.1</version>
+        </dependency>
+
         <!-- JUnit 5 for testing -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>
@@ -100,9 +112,19 @@
                             <goal>shade</goal>
                         </goals>
                         <configuration>
+                            <relocations>
+                                <relocation>
+                                    <pattern>org.objectweb.asm</pattern>
+                                    <shadedPattern>com.codeflash.asm</shadedPattern>
+                                </relocation>
+                            </relocations>
                             <transformers>
                                 <transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
                                     <mainClass>com.codeflash.Comparator</mainClass>
+                                    <manifestEntries>
+                                        <Premain-Class>com.codeflash.profiler.ProfilerAgent</Premain-Class>
+                                        <Can-Retransform-Classes>true</Can-Retransform-Classes>
+                                    </manifestEntries>
                                 </transformer>
                             </transformers>
                             <filters>

codeflash-java-runtime/src/main/java/com/codeflash/profiler/LineProfilingClassVisitor.java

@@ -0,0 +1,41 @@
+package com.codeflash.profiler;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+/**
+ * ASM ClassVisitor that filters methods and wraps target methods with
+ * {@link LineProfilingMethodVisitor} for line-level profiling.
+ */
+public class LineProfilingClassVisitor extends ClassVisitor {
+
+    private final String internalClassName;
+    private final ProfilerConfig config;
+    private String sourceFile;
+
+    public LineProfilingClassVisitor(ClassVisitor classVisitor, String internalClassName, ProfilerConfig config) {
+        super(Opcodes.ASM9, classVisitor);
+        this.internalClassName = internalClassName;
+        this.config = config;
+    }
+
+    @Override
+    public void visitSource(String source, String debug) {
+        super.visitSource(source, debug);
+        // Resolve the absolute source file path from the config
+        this.sourceFile = config.resolveSourceFile(internalClassName);
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String descriptor,
+                                     String signature, String[] exceptions) {
+        MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+
+        if (config.shouldInstrumentMethod(internalClassName, name)) {
+            return new LineProfilingMethodVisitor(mv, access, name, descriptor,
+                    internalClassName, sourceFile);
+        }
+        return mv;
+    }
+}

codeflash-java-runtime/src/main/java/com/codeflash/profiler/LineProfilingMethodVisitor.java

@@ -0,0 +1,154 @@
+package com.codeflash.profiler;
+
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.commons.AdviceAdapter;
+
+/**
+ * ASM MethodVisitor that injects line-level profiling probes.
+ *
+ * <p>At each {@code LineNumber} table entry within the target method:
+ * <ol>
+ *   <li>Registers the line with {@link ProfilerRegistry} (happens once at class-load time)</li>
+ *   <li>Injects bytecode: {@code LDC globalId; INVOKESTATIC ProfilerData.hit(I)V}</li>
+ * </ol>
+ *
+ * <p>At method entry: injects a warmup self-call loop (if warmup is configured) followed by
+ * {@code ProfilerData.enterMethod(entryLineId)}.
+ * <p>At method exit (every RETURN/ATHROW): injects {@code ProfilerData.exitMethod()}.
+ */
+public class LineProfilingMethodVisitor extends AdviceAdapter {
+
+    private static final String PROFILER_DATA = "com/codeflash/profiler/ProfilerData";
+
+    private final String internalClassName;
+    private final String sourceFile;
+    private final String methodName;
+    private boolean firstLineVisited = false;
+
+    protected LineProfilingMethodVisitor(
+            MethodVisitor mv, int access, String name, String descriptor,
+            String internalClassName, String sourceFile) {
+        super(Opcodes.ASM9, mv, access, name, descriptor);
+        this.internalClassName = internalClassName;
+        this.sourceFile = sourceFile;
+        this.methodName = name;
+    }
+
+    /**
+     * Inject a warmup self-call loop at method entry.
+     *
+     * <p>Generated bytecode equivalent:
+     * <pre>
+     * if (ProfilerData.isWarmupNeeded()) {
+     *     ProfilerData.startWarmup();
+     *     for (int i = 0; i &lt; ProfilerData.getWarmupThreshold(); i++) {
+     *         thisMethod(originalArgs);
+     *     }
+     *     ProfilerData.finishWarmup();
+     * }
+     * </pre>
+     *
+     * <p>Recursive warmup calls re-enter this method but {@code isWarmupNeeded()} returns
+     * {@code false} (guard flag set by {@code startWarmup()}), so they execute the normal
+     * instrumented body. After the loop, {@code finishWarmup()} zeros all counters so the
+     * next real execution records clean data.
+     */
+    @Override
+    protected void onMethodEnter() {
+        Label skipWarmup = new Label();
+
+        // if (!ProfilerData.isWarmupNeeded()) goto skipWarmup
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "isWarmupNeeded", "()Z", false);
+        mv.visitJumpInsn(IFEQ, skipWarmup);
+
+        // ProfilerData.startWarmup()
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "startWarmup", "()V", false);
+
+        // int _warmupIdx = 0
+        int counterLocal = newLocal(Type.INT_TYPE);
+        mv.visitInsn(ICONST_0);
+        mv.visitVarInsn(ISTORE, counterLocal);
+
+        Label loopCheck = new Label();
+        Label loopBody = new Label();
+
+        mv.visitJumpInsn(GOTO, loopCheck);
+
+        // loop body: call self with original arguments
+        mv.visitLabel(loopBody);
+
+        boolean isStatic = (methodAccess & Opcodes.ACC_STATIC) != 0;
+        if (!isStatic) {
+            loadThis();
+        }
+        loadArgs();
+
+        int invokeOp;
+        if (isStatic) {
+            invokeOp = INVOKESTATIC;
+        } else if ((methodAccess & Opcodes.ACC_PRIVATE) != 0) {
+            invokeOp = INVOKESPECIAL;
+        } else {
+            invokeOp = INVOKEVIRTUAL;
+        }
+        mv.visitMethodInsn(invokeOp, internalClassName, methodName, methodDesc, false);
+
+        // Discard return value
+        Type returnType = Type.getReturnType(methodDesc);
+        switch (returnType.getSort()) {
+            case Type.VOID:
+                break;
+            case Type.LONG:
+            case Type.DOUBLE:
+                mv.visitInsn(POP2);
+                break;
+            default:
+                mv.visitInsn(POP);
+                break;
+        }
+
+        // _warmupIdx++
+        mv.visitIincInsn(counterLocal, 1);
+
+        // loop check: _warmupIdx < ProfilerData.getWarmupThreshold()
+        mv.visitLabel(loopCheck);
+        mv.visitVarInsn(ILOAD, counterLocal);
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "getWarmupThreshold", "()I", false);
+        mv.visitJumpInsn(IF_ICMPLT, loopBody);
+
+        // ProfilerData.finishWarmup()
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "finishWarmup", "()V", false);
+
+        mv.visitLabel(skipWarmup);
+    }
+
+    @Override
+    public void visitLineNumber(int line, Label start) {
+        super.visitLineNumber(line, start);
+
+        // Register this line and get its global ID (happens once at class-load time)
+        String dotClassName = internalClassName.replace('/', '.');
+        int globalId = ProfilerRegistry.register(sourceFile, dotClassName, methodName, line);
+
+        if (!firstLineVisited) {
+            firstLineVisited = true;
+            // Inject enterMethod call at the first line of the method
+            mv.visitLdcInsn(globalId);
+            mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "enterMethod", "(I)V", false);
+        }
+
+        // Inject: ProfilerData.hit(globalId)
+        mv.visitLdcInsn(globalId);
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "hit", "(I)V", false);
+    }
+
+    @Override
+    protected void onMethodExit(int opcode) {
+        // Before every RETURN or ATHROW, flush timing for the last line
+        // This fixes the "last line always shows 0ms" bug
+        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "exitMethod", "()V", false);
+    }
+}

codeflash-java-runtime/src/main/java/com/codeflash/profiler/LineProfilingTransformer.java

@@ -0,0 +1,46 @@
+package com.codeflash.profiler;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+
+import java.lang.instrument.ClassFileTransformer;
+import java.security.ProtectionDomain;
+
+/**
+ * {@link ClassFileTransformer} that instruments target classes with line profiling.
+ *
+ * <p>When a class matches the profiler configuration, it is run through ASM
+ * to inject {@link ProfilerData#hit(int)} calls at each line number.
+ */
+public class LineProfilingTransformer implements ClassFileTransformer {
+
+    private final ProfilerConfig config;
+
+    public LineProfilingTransformer(ProfilerConfig config) {
+        this.config = config;
+    }
+
+    @Override
+    public byte[] transform(ClassLoader loader, String className,
+                            Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
+                            byte[] classfileBuffer) {
+        if (className == null || !config.shouldInstrumentClass(className)) {
+            return null; // null = don't transform
+        }
+
+        try {
+            return instrumentClass(className, classfileBuffer);
+        } catch (Exception e) {
+            System.err.println("[codeflash-profiler] Failed to instrument " + className + ": " + e.getMessage());
+            return null;
+        }
+    }
+
+    private byte[] instrumentClass(String internalClassName, byte[] bytecode) {
+        ClassReader cr = new ClassReader(bytecode);
+        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
+        LineProfilingClassVisitor cv = new LineProfilingClassVisitor(cw, internalClassName, config);
+        cr.accept(cv, ClassReader.EXPAND_FRAMES);
+        return cw.toByteArray();
+    }
+}

codeflash-java-runtime/src/main/java/com/codeflash/profiler/ProfilerAgent.java

@@ -0,0 +1,53 @@
+package com.codeflash.profiler;
+
+import java.lang.instrument.Instrumentation;
+
+/**
+ * Java agent entry point for the CodeFlash line profiler.
+ *
+ * <p>Loaded via {@code -javaagent:codeflash-profiler-agent.jar=config=/path/to/config.json}.
+ *
+ * <p>The agent:
+ * <ol>
+ *   <li>Parses the config file specifying which classes/methods to profile</li>
+ *   <li>Registers a {@link LineProfilingTransformer} to instrument target classes at load time</li>
+ *   <li>Registers a shutdown hook to write profiling results to JSON</li>
+ * </ol>
+ */
+public class ProfilerAgent {
+
+    /**
+     * Called by the JVM before {@code main()} when the agent is loaded.
+     *
+     * @param agentArgs comma-separated key=value pairs (e.g., {@code config=/path/to/config.json})
+     * @param inst      the JVM instrumentation interface
+     */
+    public static void premain(String agentArgs, Instrumentation inst) {
+        ProfilerConfig config = ProfilerConfig.parse(agentArgs);
+
+        if (config.getTargetClasses().isEmpty()) {
+            System.err.println("[codeflash-profiler] No target classes configured, profiler inactive");
+            return;
+        }
+
+        // Pre-allocate registry with estimated capacity
+        ProfilerRegistry.initialize(config.getExpectedLineCount());
+
+        // Configure warmup phase
+        ProfilerData.setWarmupThreshold(config.getWarmupIterations());
+
+        // Register the bytecode transformer
+        inst.addTransformer(new LineProfilingTransformer(config), true);
+
+        // Register shutdown hook to write results on JVM exit
+        String outputFile = config.getOutputFile();
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+            ProfilerReporter.writeResults(outputFile, config);
+        }, "codeflash-profiler-shutdown"));
+
+        int warmup = config.getWarmupIterations();
+        String warmupMsg = warmup > 0 ? ", warmup=" + warmup + " calls" : "";
+        System.err.println("[codeflash-profiler] Agent loaded, profiling "
+                + config.getTargetClasses().size() + " class(es)" + warmupMsg);
+    }
+}