codeflash-ai/codeflash
1
0
Fork 0
Code Issues 18 Pull requests 33 Projects Releases Packages Wiki Activity Actions 4.9k

implement warmup phase for Java line profiling agent

Browse source
This commit is contained in:
HeshamHM28 2026-02-20 02:44:55 +02:00
parent 0567a78618
commit 04921024ac
6 changed files with 180 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

92
codeflash-java-runtime/src/main/java/com/codeflash/profiler/LineProfilingMethodVisitor.java
View file

@ -3,6 +3,7 @@ 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;
/**
@ -14,7 +15,8 @@ import org.objectweb.asm.commons.AdviceAdapter;
* <li>Injects bytecode: {@code LDC globalId; INVOKESTATIC ProfilerData.hit(I)V}</li>
* </ol>
*
* <p>At method entry (first line): injects {@code ProfilerData.enterMethod(entryLineId)}.
* <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 {
@ -35,6 +37,94 @@ public class LineProfilingMethodVisitor extends AdviceAdapter {
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);

7
codeflash-java-runtime/src/main/java/com/codeflash/profiler/ProfilerAgent.java
View file

@ -33,6 +33,9 @@ public class ProfilerAgent {
// 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);
@ -42,7 +45,9 @@ public class ProfilerAgent {
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)");
+ config.getTargetClasses().size() + " class(es)" + warmupMsg);
}
}

8
codeflash-java-runtime/src/main/java/com/codeflash/profiler/ProfilerConfig.java
View file

@ -21,6 +21,7 @@ import java.util.Set;
public final class ProfilerConfig {
private String outputFile = "";
private int warmupIterations = 10;
private final Map<String, List<MethodTarget>> targets = new HashMap<>();
private final Map<String, String> lineContents = new HashMap<>();
private final Set<String> targetClassNames = new HashSet<>();
@ -77,6 +78,10 @@ public final class ProfilerConfig {
return outputFile;
}
public int getWarmupIterations() {
return warmupIterations;
}
public Set<String> getTargetClasses() {
return Collections.unmodifiableSet(targetClassNames);
}
@ -150,6 +155,9 @@ public final class ProfilerConfig {
case "outputFile":
this.outputFile = readString(json, pos);
break;
case "warmupIterations":
this.warmupIterations = readInt(json, pos);
break;
case "targets":
parseTargets(json, pos);
break;

66
codeflash-java-runtime/src/main/java/com/codeflash/profiler/ProfilerData.java
View file

@ -1,5 +1,6 @@
package com.codeflash.profiler;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
@ -45,6 +46,12 @@ public final class ProfilerData {
private static final List<long[]> allHitArrays = new CopyOnWriteArrayList<>();
private static final List<long[]> allTimeArrays = new CopyOnWriteArrayList<>();
// Warmup state: the method visitor injects a self-calling warmup loop,
// warmupInProgress guards against recursive re-entry into the warmup block.
private static volatile int warmupThreshold = 0;
private static volatile boolean warmupComplete = false;
private static volatile boolean warmupInProgress = false;
private ProfilerData() {}
private static long[] registerArray(long[] arr) {
@ -57,6 +64,65 @@ public final class ProfilerData {
return arr;
}
/**
* Set the number of self-call warmup iterations before measurement begins.
* Called once from {@link ProfilerAgent#premain} before any classes are loaded.
*
* @param threshold number of warmup iterations (0 = no warmup)
*/
public static void setWarmupThreshold(int threshold) {
warmupThreshold = threshold;
warmupComplete = (threshold <= 0);
}
/**
* Check whether warmup is still needed. Called by injected bytecode at target method entry.
* Returns {@code true} only on the very first call — subsequent calls (including recursive
* warmup calls) return {@code false}.
*/
public static boolean isWarmupNeeded() {
return !warmupComplete && !warmupInProgress && warmupThreshold > 0;
}
/**
* Enter warmup phase. Sets a guard flag so recursive warmup calls skip the warmup block.
*/
public static void startWarmup() {
warmupInProgress = true;
}
/**
* Return the configured warmup iteration count.
*/
public static int getWarmupThreshold() {
return warmupThreshold;
}
/**
* End warmup: zero all profiling counters, mark warmup complete, clear the guard flag.
* The next execution of the method body is the clean measurement.
*/
public static void finishWarmup() {
resetAll();
warmupComplete = true;
warmupInProgress = false;
System.err.println("[codeflash-profiler] Warmup complete after " + warmupThreshold
+ " iterations, measurement started");
}
/**
* Reset all profiling counters across all threads.
* Called once when warmup phase completes to discard warmup data.
*/
private static void resetAll() {
for (long[] arr : allHitArrays) {
Arrays.fill(arr, 0L);
}
for (long[] arr : allTimeArrays) {
Arrays.fill(arr, 0L);
}
}
/**
* Record a hit on a profiled line. This is the HOT PATH.
*

5
codeflash/languages/java/line_profiler.py
View file

@ -23,6 +23,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
AGENT_JAR_NAME = "codeflash-runtime-1.0.0.jar"
DEFAULT_WARMUP_ITERATIONS = 100
class JavaLineProfiler:
@ -37,8 +38,9 @@ class JavaLineProfiler:
"""
def __init__(self, output_file: Path) -> None:
def __init__(self, output_file: Path, warmup_iterations: int = DEFAULT_WARMUP_ITERATIONS) -> None:
self.output_file = output_file
self.warmup_iterations = warmup_iterations
def generate_agent_config(
self,
@ -87,6 +89,7 @@ class JavaLineProfiler:
config = {
"outputFile": str(self.output_file),
"warmupIterations": self.warmup_iterations,
"targets": [
{
"className": class_name,

8
codeflash/languages/java/support.py
View file

@ -298,9 +298,10 @@ class JavaSupport(LanguageSupport):
) -> bool:
"""Prepare line profiling via the bytecode-instrumentation agent.
No source files are modified. Instead, generates a config JSON that the
Java agent uses at class-load time to know which methods to instrument.
The agent is loaded via -javaagent when the JVM starts.
Generates a config JSON that the Java agent uses at class-load time to
know which methods to instrument. The agent is loaded via -javaagent
when the JVM starts. The config includes warmup iterations so the agent
discards JIT warmup data before measurement.
Args:
func_info: Function to profile.
@ -326,6 +327,7 @@ class JavaSupport(LanguageSupport):
)
self._line_profiler_agent_arg = profiler.build_javaagent_arg(config_path)
self._line_profiler_warmup_iterations = profiler.warmup_iterations
return True
except Exception:
logger.exception("Failed to prepare line profiling for %s", func_info.function_name)