Skip to content

Commit 819a56c

Browse files
authored
Merge pull request #2058 from codeflash-ai/perf/reduce-java-tracer-e2e
perf: optimize Java tracing agent (E2E reduction + serialization + writes)
2 parents 5ee642e + b737f71 commit 819a56c

16 files changed

Lines changed: 455 additions & 158 deletions

File tree

.github/workflows/ci.yaml

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ jobs:
9595
'codeflash/languages/java/' 'codeflash/languages/base.py' \
9696
'codeflash/languages/registry.py' 'codeflash/optimization/' \
9797
'codeflash/verification/' 'codeflash-java-runtime/' \
98-
'code_to_optimize/java/' 'tests/scripts/end_to_end_test_java*'
98+
'code_to_optimize/java/' 'tests/scripts/end_to_end_test_java*' \
99+
'tests/test_languages/fixtures/java_tracer_e2e/'
99100
env:
100101
MERGE_BASE: ${{ steps.merge_base.outputs.sha }}
101102

@@ -257,7 +258,7 @@ jobs:
257258
- name: init-optimization
258259
script: end_to_end_test_init_optimization.py
259260
expected_improvement: 10
260-
environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
261+
environment: ${{ ((github.event_name == 'workflow_dispatch' && github.actor != 'misrasaurabh1' && github.actor != 'KRRT7') || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
261262
runs-on: ubuntu-latest
262263
env:
263264
CODEFLASH_AIS_SERVER: prod
@@ -344,7 +345,7 @@ jobs:
344345
script: end_to_end_test_js_ts_class.py
345346
js_project_dir: code_to_optimize/js/code_to_optimize_ts
346347
expected_improvement: 30
347-
environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
348+
environment: ${{ ((github.event_name == 'workflow_dispatch' && github.actor != 'misrasaurabh1' && github.actor != 'KRRT7') || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
348349
runs-on: ubuntu-latest
349350
env:
350351
CODEFLASH_AIS_SERVER: prod
@@ -424,7 +425,7 @@ jobs:
424425
script: end_to_end_test_java_void_optimization.py
425426
expected_improvement: 70
426427
remove_git: true
427-
environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
428+
environment: ${{ ((github.event_name == 'workflow_dispatch' && github.actor != 'misrasaurabh1' && github.actor != 'KRRT7') || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
428429
runs-on: ubuntu-latest
429430
env:
430431
CODEFLASH_AIS_SERVER: prod
@@ -435,6 +436,7 @@ jobs:
435436
RETRY_DELAY: 5
436437
EXPECTED_IMPROVEMENT_PCT: ${{ matrix.expected_improvement }}
437438
CODEFLASH_END_TO_END: 1
439+
CODEFLASH_LOOPING_TIME: 5
438440
steps:
439441
- uses: actions/checkout@v6
440442
with:
@@ -468,7 +470,15 @@ jobs:
468470
- name: Install dependencies
469471
run: uv sync
470472

473+
- name: Cache codeflash-runtime JAR
474+
id: runtime-jar-cache
475+
uses: actions/cache@v4
476+
with:
477+
path: ~/.m2/repository/io/codeflash
478+
key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}
479+
471480
- name: Build and install codeflash-runtime JAR
481+
if: steps.runtime-jar-cache.outputs.cache-hit != 'true'
472482
run: |
473483
cd codeflash-java-runtime
474484
mvn install -q -DskipTests

.github/workflows/codeflash-optimize.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,4 +43,4 @@ jobs:
4343
- name: ⚡️Codeflash Optimization
4444
id: optimize_code
4545
run: |
46-
uv run codeflash --benchmark --testgen-review
46+
uv run codeflash --benchmark --testgen-review --no-pr

.github/workflows/java-e2e.yaml

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
name: Java E2E Tests
2+
on:
3+
workflow_dispatch:
4+
5+
jobs:
6+
e2e-java:
7+
strategy:
8+
fail-fast: false
9+
matrix:
10+
include:
11+
- name: java-fibonacci-nogit
12+
script: end_to_end_test_java_fibonacci.py
13+
expected_improvement: 70
14+
remove_git: true
15+
- name: java-tracer
16+
script: end_to_end_test_java_tracer.py
17+
expected_improvement: 10
18+
- name: java-void-optimization-nogit
19+
script: end_to_end_test_java_void_optimization.py
20+
expected_improvement: 70
21+
remove_git: true
22+
runs-on: ubuntu-latest
23+
env:
24+
CODEFLASH_AIS_SERVER: prod
25+
POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }}
26+
CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
27+
COLUMNS: 110
28+
MAX_RETRIES: 3
29+
RETRY_DELAY: 5
30+
EXPECTED_IMPROVEMENT_PCT: ${{ matrix.expected_improvement }}
31+
CODEFLASH_END_TO_END: 1
32+
CODEFLASH_LOOPING_TIME: 5
33+
steps:
34+
- uses: actions/checkout@v6
35+
36+
- name: Set up JDK 11
37+
uses: actions/setup-java@v5
38+
with:
39+
java-version: '11'
40+
distribution: 'temurin'
41+
cache: maven
42+
43+
- name: Install uv
44+
uses: astral-sh/setup-uv@v8.0.0
45+
with:
46+
python-version: 3.11.6
47+
enable-cache: true
48+
49+
- name: Install dependencies
50+
run: uv sync
51+
52+
- name: Cache codeflash-runtime JAR
53+
id: runtime-jar-cache
54+
uses: actions/cache@v4
55+
with:
56+
path: ~/.m2/repository/io/codeflash
57+
key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}
58+
59+
- name: Build and install codeflash-runtime JAR
60+
if: steps.runtime-jar-cache.outputs.cache-hit != 'true'
61+
run: |
62+
cd codeflash-java-runtime
63+
mvn install -q -DskipTests
64+
65+
- name: Remove .git
66+
if: matrix.remove_git
67+
run: |
68+
if [ -d ".git" ]; then
69+
sudo rm -rf .git
70+
echo ".git directory removed."
71+
else
72+
echo ".git directory does not exist."
73+
exit 1
74+
fi
75+
76+
- name: Run E2E test
77+
run: uv run python tests/scripts/${{ matrix.script }}

codeflash-java-runtime/src/main/java/com/codeflash/Serializer.java

Lines changed: 101 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import com.esotericsoftware.kryo.util.DefaultInstantiatorStrategy;
77
import org.objenesis.strategy.StdInstantiatorStrategy;
88

9-
import java.io.ByteArrayOutputStream;
109
import java.io.InputStream;
1110
import java.io.OutputStream;
1211
import java.lang.reflect.Field;
@@ -36,7 +35,11 @@ public final class Serializer {
3635
private static final int MAX_COLLECTION_SIZE = 1000;
3736
private static final int BUFFER_SIZE = 4096;
3837

39-
// Thread-local Kryo instances (Kryo is not thread-safe)
38+
// Thread-local Kryo, Output, and IdentityHashMap instances for reuse
39+
private static final ThreadLocal<Output> OUTPUT = ThreadLocal.withInitial(() -> new Output(BUFFER_SIZE, -1));
40+
private static final ThreadLocal<IdentityHashMap<Object, Object>> SEEN =
41+
ThreadLocal.withInitial(IdentityHashMap::new);
42+
4043
private static final ThreadLocal<Kryo> KRYO = ThreadLocal.withInitial(() -> {
4144
Kryo kryo = new Kryo();
4245
kryo.setRegistrationRequired(false);
@@ -89,10 +92,78 @@ private Serializer() {
8992
* @return Serialized bytes (may contain KryoPlaceholder for unserializable parts)
9093
*/
9194
public static byte[] serialize(Object obj) {
92-
Object processed = recursiveProcess(obj, new IdentityHashMap<>(), 0, "");
95+
// Fast path: if args are all safe types, skip recursive processing entirely
96+
if (obj instanceof Object[] && isSafeArgs((Object[]) obj)) {
97+
return directSerialize(obj);
98+
}
99+
100+
IdentityHashMap<Object, Object> seen = SEEN.get();
101+
seen.clear();
102+
Object processed = recursiveProcess(obj, seen, 0, "");
93103
return directSerialize(processed);
94104
}
95105

106+
/**
107+
* Attempt fast-path serialization for args that are all known-safe types.
108+
* Returns serialized bytes if all args are safe, or null if the slow path is needed.
109+
* Callers can use this to avoid executor submission overhead for simple arguments.
110+
*/
111+
public static byte[] serializeFast(Object obj) {
112+
if (obj instanceof Object[] && isSafeArgs((Object[]) obj)) {
113+
return directSerialize(obj);
114+
}
115+
return null;
116+
}
117+
118+
/**
119+
* Check if all elements of an args array can be serialized directly without recursive processing.
120+
*/
121+
private static boolean isSafeArgs(Object[] args) {
122+
for (Object arg : args) {
123+
if (!isSafeForDirectSerialization(arg)) {
124+
return false;
125+
}
126+
}
127+
return true;
128+
}
129+
130+
/**
131+
* Check if an object is safe to serialize directly without recursive processing.
132+
* Covers: null, simple types, primitive arrays, and safe containers (up to 3 levels deep).
133+
*/
134+
private static boolean isSafeForDirectSerialization(Object obj) {
135+
return isSafeForDirectSerialization(obj, 3);
136+
}
137+
138+
private static boolean isSafeForDirectSerialization(Object obj, int depthLeft) {
139+
if (obj == null || isSimpleType(obj)) {
140+
return true;
141+
}
142+
if (depthLeft <= 0) {
143+
return false;
144+
}
145+
Class<?> clazz = obj.getClass();
146+
if (clazz.isArray() && clazz.getComponentType().isPrimitive()) {
147+
return true;
148+
}
149+
if (isSafeContainerType(clazz)) {
150+
if (obj instanceof Collection) {
151+
for (Object item : (Collection<?>) obj) {
152+
if (!isSafeForDirectSerialization(item, depthLeft - 1)) return false;
153+
}
154+
return true;
155+
}
156+
if (obj instanceof Map) {
157+
for (Map.Entry<?, ?> e : ((Map<?, ?>) obj).entrySet()) {
158+
if (!isSafeForDirectSerialization(e.getKey(), depthLeft - 1) ||
159+
!isSafeForDirectSerialization(e.getValue(), depthLeft - 1)) return false;
160+
}
161+
return true;
162+
}
163+
}
164+
return false;
165+
}
166+
96167
/**
97168
* Deserialize bytes back to an object.
98169
* The returned object may contain KryoPlaceholder instances for parts
@@ -141,14 +212,15 @@ public static byte[] serializeException(Throwable error) {
141212

142213
/**
143214
* Direct serialization without recursive processing.
215+
* Reuses a ThreadLocal Output buffer to avoid per-call allocation.
144216
*/
145217
private static byte[] directSerialize(Object obj) {
146218
Kryo kryo = KRYO.get();
147-
ByteArrayOutputStream baos = new ByteArrayOutputStream(BUFFER_SIZE);
148-
try (Output output = new Output(baos)) {
149-
kryo.writeClassAndObject(output, obj);
150-
}
151-
return baos.toByteArray();
219+
Output output = OUTPUT.get();
220+
output.reset();
221+
kryo.writeClassAndObject(output, obj);
222+
output.flush();
223+
return output.toBytes();
152224
}
153225

154226
/**
@@ -201,37 +273,23 @@ private static Object recursiveProcess(Object obj, IdentityHashMap<Object, Objec
201273
// unserializable types, recursively process to catch and replace unserializable objects.
202274
if (obj instanceof Map) {
203275
Map<?, ?> map = (Map<?, ?>) obj;
204-
if (containsOnlySimpleTypes(map)) {
205-
// Simple map - try direct serialization to preserve full size
206-
byte[] serialized = tryDirectSerialize(obj);
207-
if (serialized != null) {
208-
try {
209-
deserialize(serialized);
210-
return obj; // Success - return original
211-
} catch (Exception e) {
212-
// Fall through to recursive handling
213-
}
214-
}
276+
if (isSafeContainerType(clazz) && containsOnlySimpleTypes(map)) {
277+
return obj;
215278
}
216279
return handleMap(map, seen, depth, path);
217280
}
218281
if (obj instanceof Collection) {
219282
Collection<?> collection = (Collection<?>) obj;
220-
if (containsOnlySimpleTypes(collection)) {
221-
// Simple collection - try direct serialization to preserve full size
222-
byte[] serialized = tryDirectSerialize(obj);
223-
if (serialized != null) {
224-
try {
225-
deserialize(serialized);
226-
return obj; // Success - return original
227-
} catch (Exception e) {
228-
// Fall through to recursive handling
229-
}
230-
}
283+
if (isSafeContainerType(clazz) && containsOnlySimpleTypes(collection)) {
284+
return obj;
231285
}
232286
return handleCollection(collection, seen, depth, path);
233287
}
234288
if (clazz.isArray()) {
289+
// Primitive arrays (int[], double[], etc.) are directly serializable by Kryo
290+
if (clazz.getComponentType().isPrimitive()) {
291+
return obj;
292+
}
235293
return handleArray(obj, seen, depth, path);
236294
}
237295

@@ -255,6 +313,19 @@ private static Object recursiveProcess(Object obj, IdentityHashMap<Object, Objec
255313
}
256314
}
257315

316+
/**
317+
* Check if a container type is known to round-trip safely through Kryo without verification.
318+
* Only includes types registered with Kryo that are known to serialize/deserialize correctly.
319+
*/
320+
private static boolean isSafeContainerType(Class<?> clazz) {
321+
return clazz == ArrayList.class ||
322+
clazz == LinkedList.class ||
323+
clazz == HashMap.class ||
324+
clazz == LinkedHashMap.class ||
325+
clazz == HashSet.class ||
326+
clazz == LinkedHashSet.class;
327+
}
328+
258329
/**
259330
* Check if a class is known to be unserializable.
260331
*/

codeflash-java-runtime/src/main/java/com/codeflash/tracer/TraceRecorder.java

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ public final class TraceRecorder {
3131

3232
private TraceRecorder(TracerConfig config) {
3333
this.config = config;
34-
this.writer = new TraceWriter(config.getDbPath());
34+
this.writer = new TraceWriter(config.getDbPath(), config.isInMemoryDb());
3535
this.maxFunctionCount = config.getMaxFunctionCount();
3636
this.serializerExecutor = Executors.newCachedThreadPool(r -> {
3737
Thread t = new Thread(r, "codeflash-serializer");
@@ -76,23 +76,27 @@ private void onEntryImpl(String className, String methodName, String descriptor,
7676
return;
7777
}
7878

79-
// Serialize args with timeout to prevent deep object graph traversal from blocking
79+
// Serialize args — try inline fast path first, fall back to async with timeout
8080
byte[] argsBlob;
81-
Future<byte[]> future = serializerExecutor.submit(() -> Serializer.serialize(args));
82-
try {
83-
argsBlob = future.get(SERIALIZATION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
84-
} catch (TimeoutException e) {
85-
future.cancel(true);
86-
droppedCaptures.incrementAndGet();
87-
System.err.println("[codeflash-tracer] Serialization timed out for " + className + "."
88-
+ methodName);
89-
return;
90-
} catch (Exception e) {
91-
Throwable cause = e.getCause() != null ? e.getCause() : e;
92-
droppedCaptures.incrementAndGet();
93-
System.err.println("[codeflash-tracer] Serialization failed for " + className + "."
94-
+ methodName + ": " + cause.getClass().getSimpleName() + ": " + cause.getMessage());
95-
return;
81+
argsBlob = Serializer.serializeFast(args);
82+
if (argsBlob == null) {
83+
// Slow path: async serialization with timeout for complex/unknown types
84+
Future<byte[]> future = serializerExecutor.submit(() -> Serializer.serialize(args));
85+
try {
86+
argsBlob = future.get(SERIALIZATION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
87+
} catch (TimeoutException e) {
88+
future.cancel(true);
89+
droppedCaptures.incrementAndGet();
90+
System.err.println("[codeflash-tracer] Serialization timed out for " + className + "."
91+
+ methodName);
92+
return;
93+
} catch (Exception e) {
94+
Throwable cause = e.getCause() != null ? e.getCause() : e;
95+
droppedCaptures.incrementAndGet();
96+
System.err.println("[codeflash-tracer] Serialization failed for " + className + "."
97+
+ methodName + ": " + cause.getClass().getSimpleName() + ": " + cause.getMessage());
98+
return;
99+
}
96100
}
97101

98102
long timeNs = System.nanoTime();

0 commit comments

Comments
 (0)