Try to make diff generation more robust (cc #19)

makamys · Oct 24, 2023 · 95d072f · 95d072f
1 parent 8c1e56d
commit 95d072f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 13 deletions.
diff --git a/.../java/makamys/coretweaks/optimization/transformercache/lite/CachedTransformerWrapper.java b/.../java/makamys/coretweaks/optimization/transformercache/lite/CachedTransformerWrapper.java
@@ -22,7 +22,9 @@ public byte[] wrapTransform(String name, String transformedName, byte[] basicCla
             misses++;
             TransformerCache.instance.prePutCached(transformerName, name, transformedName, basicClass);
             result = proxy.invokeNextHandler(name, transformedName, basicClass);
-            TransformerCache.instance.putCached(transformerName, name, transformedName, result);
+            if(!TransformerCache.instance.putCached(transformerName, name, transformedName, result)) {
+                return basicClass;
+            }
         }
         return TransformerCache.fromNullableByteArray(result);
     }

diff --git a/src/main/java/makamys/coretweaks/optimization/transformercache/lite/TransformerCache.java b/src/main/java/makamys/coretweaks/optimization/transformercache/lite/TransformerCache.java
@@ -12,6 +12,7 @@
 import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.channels.ClosedByInterruptException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -288,12 +289,17 @@ private void putLastClassData(byte[] data) {
     }
 
     /** MUST be preceded with a call to prePutCached. */
-    public void putCached(String transName, String name, String transformedName, byte[] result) {
+    public boolean putCached(String transName, String name, String transformedName, byte[] result) {
         TransformerData data = transformerMap.get(transName);
         if(data == null) {
             transformerMap.put(transName, data = new TransformerData(transName));
         }
-        data.transformationMap.put(transformedName, new CachedTransformation(transformedName, lastClassData, lastClassDataLength, result));
+        CachedTransformation cached = new CachedTransformation(transformedName, lastClassData, lastClassDataLength, result);
+        if(cached.isValid()) {
+            data.transformationMap.put(transformedName, cached);
+            return true;
+        }
+        return false;
     }
 
     public static int calculateHash(byte[] data) {
@@ -330,6 +336,8 @@ public TransformerData(String transformerClassName) {
         public TransformerData() {}
 
         public static class CachedTransformation {
+            private static final byte[] INVALID_RESULT = new byte[] {};
+
             String targetClassName;
             int preLength;
             int preHash;
@@ -340,32 +348,45 @@ public static class CachedTransformation {
 
             public CachedTransformation() {}
 
+            public boolean isValid() {
+                return diff != INVALID_RESULT;
+            }
+
             public CachedTransformation(String targetClassName, byte[] source, int sourceLen, byte[] target) {
                 this.targetClassName = targetClassName;
                 this.preHash = calculateHash(source, sourceLen);
                 this.preLength = sourceLen;
                 this.postLength = nullSafeLength(target);
                 this.postHash = calculateHash(target);
                 if(preHash != postHash) {
-                    diff = generateDiff(source, sourceLen, target);
+                    diff = generateDiff(source, sourceLen, target, targetClassName);
                 }
                 this.lastAccessed = now();
             }
 
-            private static byte[] generateDiff(byte[] source, int sourceLen, byte[] target) {
+            private static byte[] generateDiff(byte[] source, int sourceLen, byte[] target, String name) {
                 if(source == null || !TransformerCache.instance.meta.enableDiffs) {
                     return target;
                 }
 
-                try {
-                    ByteArrayOutputStream os = new ByteArrayOutputStream();
-                    delta.compute(new FastByteBufferSeekableSource(ByteBuffer.wrap(source, 0, sourceLen)), new ByteArrayInputStream(target), new GDiffWriter(os));
-                    return os.toByteArray();
-                } catch(Exception e) {
-                    LOGGER.error("Failed to generate diff");
-                    e.printStackTrace();
-                    throw new RuntimeException(e);
+                for(int attemptsLeft = 4; attemptsLeft >= 0; attemptsLeft--) {
+                    try {
+                        ByteArrayOutputStream os = new ByteArrayOutputStream();
+                        delta.compute(new FastByteBufferSeekableSource(ByteBuffer.wrap(source, 0, sourceLen)), new ByteArrayInputStream(target), new GDiffWriter(os));
+                        return os.toByteArray();
+                    } catch(Exception e) {
+                        if(!(e instanceof ClosedByInterruptException)) {
+                            attemptsLeft = 0;
+                        }
+                        if(attemptsLeft > 0) {
+                            LOGGER.error("Failed to generate diff for class " + name + ", will try again " + attemptsLeft + " more times");
+                        } else {
+                            LOGGER.error("Failed to generate diff for class " + name + ". Please report this if it keeps happening!");
+                        }
+                        e.printStackTrace();
+                    }
                 }
+                return INVALID_RESULT;
             }
 
             public byte[] getNewClass(byte[] source) {