perf: optimize Java tracing agent (E2E reduction + serialization + writes)

.github/workflows/ci.yaml

@@ -95,7 +95,8 @@ jobs:
             'codeflash/languages/java/' 'codeflash/languages/base.py' \
             'codeflash/languages/registry.py' 'codeflash/optimization/' \
             'codeflash/verification/' 'codeflash-java-runtime/' \
-            'code_to_optimize/java/' 'tests/scripts/end_to_end_test_java*'
+            'code_to_optimize/java/' 'tests/scripts/end_to_end_test_java*' \
+            'tests/test_languages/fixtures/java_tracer_e2e/'
         env:
           MERGE_BASE: ${{ steps.merge_base.outputs.sha }}
 
@@ -257,7 +258,7 @@ jobs:
           - name: init-optimization
             script: end_to_end_test_init_optimization.py
             expected_improvement: 10
-    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' || '' }}
+    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' || '' }}
     runs-on: ubuntu-latest
     env:
       CODEFLASH_AIS_SERVER: prod
@@ -344,7 +345,7 @@ jobs:
             script: end_to_end_test_js_ts_class.py
             js_project_dir: code_to_optimize/js/code_to_optimize_ts
             expected_improvement: 30
-    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' || '' }}
+    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' || '' }}
     runs-on: ubuntu-latest
     env:
       CODEFLASH_AIS_SERVER: prod
@@ -424,7 +425,7 @@ jobs:
             script: end_to_end_test_java_void_optimization.py
             expected_improvement: 70
             remove_git: true
-    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' || '' }}
+    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' || '' }}
     runs-on: ubuntu-latest
     env:
       CODEFLASH_AIS_SERVER: prod
@@ -435,6 +436,7 @@ jobs:
       RETRY_DELAY: 5
       EXPECTED_IMPROVEMENT_PCT: ${{ matrix.expected_improvement }}
       CODEFLASH_END_TO_END: 1
+      CODEFLASH_LOOPING_TIME: 5
     steps:
       - uses: actions/checkout@v6
         with:
@@ -468,7 +470,15 @@ jobs:
       - name: Install dependencies
         run: uv sync
 
+      - name: Cache codeflash-runtime JAR
+        id: runtime-jar-cache
+        uses: actions/cache@v4
+        with:
+          path: ~/.m2/repository/io/codeflash
+          key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}
+
       - name: Build and install codeflash-runtime JAR
+        if: steps.runtime-jar-cache.outputs.cache-hit != 'true'
         run: |
           cd codeflash-java-runtime
           mvn install -q -DskipTests

.github/workflows/codeflash-optimize.yaml

@@ -43,4 +43,4 @@ jobs:
       - name: ⚡️Codeflash Optimization
         id: optimize_code
         run: |
-          uv run codeflash --benchmark --testgen-review
+          uv run codeflash --benchmark --testgen-review --no-pr

.github/workflows/java-e2e.yaml

@@ -0,0 +1,77 @@
+name: Java E2E Tests
+on:
+  workflow_dispatch:
+
+jobs:
+  e2e-java:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: java-fibonacci-nogit
+            script: end_to_end_test_java_fibonacci.py
+            expected_improvement: 70
+            remove_git: true
+          - name: java-tracer
+            script: end_to_end_test_java_tracer.py
+            expected_improvement: 10
+          - name: java-void-optimization-nogit
+            script: end_to_end_test_java_void_optimization.py
+            expected_improvement: 70
+            remove_git: true
+    runs-on: ubuntu-latest
+    env:
+      CODEFLASH_AIS_SERVER: prod
+      POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }}
+      CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
+      COLUMNS: 110
+      MAX_RETRIES: 3
+      RETRY_DELAY: 5
+      EXPECTED_IMPROVEMENT_PCT: ${{ matrix.expected_improvement }}
+      CODEFLASH_END_TO_END: 1
+      CODEFLASH_LOOPING_TIME: 5
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up JDK 11
+        uses: actions/setup-java@v5
+        with:
+          java-version: '11'
+          distribution: 'temurin'
+          cache: maven
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v8.0.0
+        with:
+          python-version: 3.11.6
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Cache codeflash-runtime JAR
+        id: runtime-jar-cache
+        uses: actions/cache@v4
+        with:
+          path: ~/.m2/repository/io/codeflash
+          key: codeflash-runtime-${{ hashFiles('codeflash-java-runtime/pom.xml', 'codeflash-java-runtime/src/**') }}
+
+      - name: Build and install codeflash-runtime JAR
+        if: steps.runtime-jar-cache.outputs.cache-hit != 'true'
+        run: |
+          cd codeflash-java-runtime
+          mvn install -q -DskipTests
+
+      - name: Remove .git
+        if: matrix.remove_git
+        run: |
+          if [ -d ".git" ]; then
+            sudo rm -rf .git
+            echo ".git directory removed."
+          else
+            echo ".git directory does not exist."
+            exit 1
+          fi
+
+      - name: Run E2E test
+        run: uv run python tests/scripts/${{ matrix.script }}

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

@@ -6,7 +6,6 @@
 import com.esotericsoftware.kryo.util.DefaultInstantiatorStrategy;
 import org.objenesis.strategy.StdInstantiatorStrategy;
 
-import java.io.ByteArrayOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.reflect.Field;
@@ -36,7 +35,11 @@ public final class Serializer {
     private static final int MAX_COLLECTION_SIZE = 1000;
     private static final int BUFFER_SIZE = 4096;
 
-    // Thread-local Kryo instances (Kryo is not thread-safe)
+    // Thread-local Kryo, Output, and IdentityHashMap instances for reuse
+    private static final ThreadLocal<Output> OUTPUT = ThreadLocal.withInitial(() -> new Output(BUFFER_SIZE, -1));
+    private static final ThreadLocal<IdentityHashMap<Object, Object>> SEEN =
+        ThreadLocal.withInitial(IdentityHashMap::new);
+
     private static final ThreadLocal<Kryo> KRYO = ThreadLocal.withInitial(() -> {
         Kryo kryo = new Kryo();
         kryo.setRegistrationRequired(false);
@@ -89,10 +92,78 @@ private Serializer() {
      * @return Serialized bytes (may contain KryoPlaceholder for unserializable parts)
      */
     public static byte[] serialize(Object obj) {
-        Object processed = recursiveProcess(obj, new IdentityHashMap<>(), 0, "");
+        // Fast path: if args are all safe types, skip recursive processing entirely
+        if (obj instanceof Object[] && isSafeArgs((Object[]) obj)) {
+            return directSerialize(obj);
+        }
+
+        IdentityHashMap<Object, Object> seen = SEEN.get();
+        seen.clear();
+        Object processed = recursiveProcess(obj, seen, 0, "");
         return directSerialize(processed);
     }
 
+    /**
+     * Attempt fast-path serialization for args that are all known-safe types.
+     * Returns serialized bytes if all args are safe, or null if the slow path is needed.
+     * Callers can use this to avoid executor submission overhead for simple arguments.
+     */
+    public static byte[] serializeFast(Object obj) {
+        if (obj instanceof Object[] && isSafeArgs((Object[]) obj)) {
+            return directSerialize(obj);
+        }
+        return null;
+    }
+
+    /**
+     * Check if all elements of an args array can be serialized directly without recursive processing.
+     */
+    private static boolean isSafeArgs(Object[] args) {
+        for (Object arg : args) {
+            if (!isSafeForDirectSerialization(arg)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Check if an object is safe to serialize directly without recursive processing.
+     * Covers: null, simple types, primitive arrays, and safe containers (up to 3 levels deep).
+     */
+    private static boolean isSafeForDirectSerialization(Object obj) {
+        return isSafeForDirectSerialization(obj, 3);
+    }
+
+    private static boolean isSafeForDirectSerialization(Object obj, int depthLeft) {
+        if (obj == null || isSimpleType(obj)) {
+            return true;
+        }
+        if (depthLeft <= 0) {
+            return false;
+        }
+        Class<?> clazz = obj.getClass();
+        if (clazz.isArray() && clazz.getComponentType().isPrimitive()) {
+            return true;
+        }
+        if (isSafeContainerType(clazz)) {
+            if (obj instanceof Collection) {
+                for (Object item : (Collection<?>) obj) {
+                    if (!isSafeForDirectSerialization(item, depthLeft - 1)) return false;
+                }
+                return true;
+            }
+            if (obj instanceof Map) {
+                for (Map.Entry<?, ?> e : ((Map<?, ?>) obj).entrySet()) {
+                    if (!isSafeForDirectSerialization(e.getKey(), depthLeft - 1) ||
+                        !isSafeForDirectSerialization(e.getValue(), depthLeft - 1)) return false;
+                }
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Deserialize bytes back to an object.
      * The returned object may contain KryoPlaceholder instances for parts
@@ -141,14 +212,15 @@ public static byte[] serializeException(Throwable error) {
 
     /**
      * Direct serialization without recursive processing.
+     * Reuses a ThreadLocal Output buffer to avoid per-call allocation.
      */
     private static byte[] directSerialize(Object obj) {
         Kryo kryo = KRYO.get();
-        ByteArrayOutputStream baos = new ByteArrayOutputStream(BUFFER_SIZE);
-        try (Output output = new Output(baos)) {
-            kryo.writeClassAndObject(output, obj);
-        }
-        return baos.toByteArray();
+        Output output = OUTPUT.get();
+        output.reset();
+        kryo.writeClassAndObject(output, obj);
+        output.flush();
+        return output.toBytes();
     }
 
     /**
@@ -201,37 +273,23 @@ private static Object recursiveProcess(Object obj, IdentityHashMap<Object, Objec
             // unserializable types, recursively process to catch and replace unserializable objects.
             if (obj instanceof Map) {
                 Map<?, ?> map = (Map<?, ?>) obj;
-                if (containsOnlySimpleTypes(map)) {
-                    // Simple map - try direct serialization to preserve full size
-                    byte[] serialized = tryDirectSerialize(obj);
-                    if (serialized != null) {
-                        try {
-                            deserialize(serialized);
-                            return obj; // Success - return original
-                        } catch (Exception e) {
-                            // Fall through to recursive handling
-                        }
-                    }
+                if (isSafeContainerType(clazz) && containsOnlySimpleTypes(map)) {
+                    return obj;
                 }
                 return handleMap(map, seen, depth, path);
             }
             if (obj instanceof Collection) {
                 Collection<?> collection = (Collection<?>) obj;
-                if (containsOnlySimpleTypes(collection)) {
-                    // Simple collection - try direct serialization to preserve full size
-                    byte[] serialized = tryDirectSerialize(obj);
-                    if (serialized != null) {
-                        try {
-                            deserialize(serialized);
-                            return obj; // Success - return original
-                        } catch (Exception e) {
-                            // Fall through to recursive handling
-                        }
-                    }
+                if (isSafeContainerType(clazz) && containsOnlySimpleTypes(collection)) {
+                    return obj;
                 }
                 return handleCollection(collection, seen, depth, path);
             }
             if (clazz.isArray()) {
+                // Primitive arrays (int[], double[], etc.) are directly serializable by Kryo
+                if (clazz.getComponentType().isPrimitive()) {
+                    return obj;
+                }
                 return handleArray(obj, seen, depth, path);
             }
 
@@ -255,6 +313,19 @@ private static Object recursiveProcess(Object obj, IdentityHashMap<Object, Objec
         }
     }
 
+    /**
+     * Check if a container type is known to round-trip safely through Kryo without verification.
+     * Only includes types registered with Kryo that are known to serialize/deserialize correctly.
+     */
+    private static boolean isSafeContainerType(Class<?> clazz) {
+        return clazz == ArrayList.class ||
+               clazz == LinkedList.class ||
+               clazz == HashMap.class ||
+               clazz == LinkedHashMap.class ||
+               clazz == HashSet.class ||
+               clazz == LinkedHashSet.class;
+    }
+
     /**
      * Check if a class is known to be unserializable.
      */

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

@@ -31,7 +31,7 @@ public final class TraceRecorder {
 
     private TraceRecorder(TracerConfig config) {
         this.config = config;
-        this.writer = new TraceWriter(config.getDbPath());
+        this.writer = new TraceWriter(config.getDbPath(), config.isInMemoryDb());
         this.maxFunctionCount = config.getMaxFunctionCount();
         this.serializerExecutor = Executors.newCachedThreadPool(r -> {
             Thread t = new Thread(r, "codeflash-serializer");
@@ -76,23 +76,27 @@ private void onEntryImpl(String className, String methodName, String descriptor,
             return;
         }
 
-        // Serialize args with timeout to prevent deep object graph traversal from blocking
+        // Serialize args — try inline fast path first, fall back to async with timeout
         byte[] argsBlob;
-        Future<byte[]> future = serializerExecutor.submit(() -> Serializer.serialize(args));
-        try {
-            argsBlob = future.get(SERIALIZATION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
-        } catch (TimeoutException e) {
-            future.cancel(true);
-            droppedCaptures.incrementAndGet();
-            System.err.println("[codeflash-tracer] Serialization timed out for " + className + "."
-                    + methodName);
-            return;
-        } catch (Exception e) {
-            Throwable cause = e.getCause() != null ? e.getCause() : e;
-            droppedCaptures.incrementAndGet();
-            System.err.println("[codeflash-tracer] Serialization failed for " + className + "."
-                    + methodName + ": " + cause.getClass().getSimpleName() + ": " + cause.getMessage());
-            return;
+        argsBlob = Serializer.serializeFast(args);
+        if (argsBlob == null) {
+            // Slow path: async serialization with timeout for complex/unknown types
+            Future<byte[]> future = serializerExecutor.submit(() -> Serializer.serialize(args));
+            try {
+                argsBlob = future.get(SERIALIZATION_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            } catch (TimeoutException e) {
+                future.cancel(true);
+                droppedCaptures.incrementAndGet();
+                System.err.println("[codeflash-tracer] Serialization timed out for " + className + "."
+                        + methodName);
+                return;
+            } catch (Exception e) {
+                Throwable cause = e.getCause() != null ? e.getCause() : e;
+                droppedCaptures.incrementAndGet();
+                System.err.println("[codeflash-tracer] Serialization failed for " + className + "."
+                        + methodName + ": " + cause.getClass().getSimpleName() + ": " + cause.getMessage());
+                return;
+            }
         }
 
         long timeNs = System.nanoTime();