Fixed buffer corruption issues if the buffer is used over multiple frames

Fixes #298

Author: RaphiMC

common/src/main/java/net/raphimc/immediatelyfast/feature/core/BatchableBufferSource.java

@@ -70,14 +70,14 @@ public VertexConsumer getBuffer(final RenderLayer layer) {
         final BufferBuilder bufferBuilder;
         boolean hasBufferForRenderLayer = layer.areVerticesNotShared() && this.pendingBuffers.containsKey(layer);
         if (!layer.areVerticesNotShared()) {
-            bufferBuilder = new BufferBuilder(BufferAllocatorPool.borrowBufferAllocator(), layer.getDrawMode(), layer.getVertexFormat());
+            bufferBuilder = new BufferBuilder(this.getNextBufferAllocator(), layer.getDrawMode(), layer.getVertexFormat());
             this.currentLayer = layer;
         } else if (hasBufferForRenderLayer) {
             bufferBuilder = this.pendingBuffers.get(layer).iterator().next();
         } else if (this.layerBuffers.containsKey(layer)) {
             bufferBuilder = new BufferBuilder(this.layerBuffers.get(layer), layer.getDrawMode(), layer.getVertexFormat());
         } else {
-            bufferBuilder = new BufferBuilder(BufferAllocatorPool.borrowBufferAllocator(), layer.getDrawMode(), layer.getVertexFormat());
+            bufferBuilder = new BufferBuilder(this.getNextBufferAllocator(), layer.getDrawMode(), layer.getVertexFormat());
             this.currentLayer = layer;
         }
 
@@ -168,13 +168,16 @@ public void drawDirect(final RenderLayer layer) {
 
         this.activeLayers.remove(layer);
         for (BufferBuilder bufferBuilder : this.getBufferBuilder(layer)) {
-            final BufferAllocator prevBufferAllocator = bufferBuilder.allocator;
+            final BufferAllocator prevBufferAllocator = this.allocator;
             this.allocator = bufferBuilder.allocator;
             this.draw(layer, bufferBuilder);
             this.allocator = prevBufferAllocator;
             BufferAllocatorPool.returnBufferAllocatorSafe(bufferBuilder.allocator);
         }
         this.pendingBuffers.remove(layer);
+        if (this.currentLayer == layer) {
+            this.currentLayer = null;
+        }
 
         if (IrisCompat.IRIS_LOADED && !IrisCompat.isRenderingLevel.getAsBoolean()) {
             IrisCompat.renderWithExtendedVertexFormat.accept(true);
@@ -225,4 +228,12 @@ protected int getLayerOrder(final RenderLayer layer) {
         }
     }
 
+    private BufferAllocator getNextBufferAllocator() {
+        if (this.allocator != FALLBACK_BUFFER && this.currentLayer == null && this.allocator.pointer != 0L) {
+            return this.allocator;
+        } else {
+            return BufferAllocatorPool.borrowBufferAllocator();
+        }
+    }
+
 }

common/src/main/java/net/raphimc/immediatelyfast/feature/core/BufferAllocatorPool.java

@@ -18,46 +18,48 @@
 package net.raphimc.immediatelyfast.feature.core;
 
 import com.mojang.blaze3d.systems.RenderSystem;
-import it.unimi.dsi.fastutil.objects.Reference2LongMap;
-import it.unimi.dsi.fastutil.objects.Reference2LongOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Reference2ObjectMap;
+import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ReferenceArrayList;
 import it.unimi.dsi.fastutil.objects.ReferenceList;
 import net.minecraft.client.util.BufferAllocator;
 import net.raphimc.immediatelyfast.ImmediatelyFast;
 
 public class BufferAllocatorPool {
 
-    private static final ReferenceList<BufferAllocator> FREE = new ReferenceArrayList<>();
-    private static final ReferenceList<BufferAllocator> IN_USE = new ReferenceArrayList<>();
-    private static final Reference2LongMap<BufferAllocator> BUFFER_ALLOCATOR_ACCESS_TIME = new Reference2LongOpenHashMap<>();
+    private static final ReferenceList<Entry> FREE = new ReferenceArrayList<>();
+    private static final ReferenceList<Entry> IN_USE = new ReferenceArrayList<>();
+    private static final Reference2ObjectMap<BufferAllocator, Entry> BUFFER_ALLOCATOR_MAPPING = new Reference2ObjectOpenHashMap<>();
 
     private BufferAllocatorPool() {
     }
 
     public static BufferAllocator borrowBufferAllocator() {
         RenderSystem.assertOnRenderThread();
-        BufferAllocator bufferAllocator;
+        Entry entry;
         if (FREE.isEmpty()) {
-            bufferAllocator = new BufferAllocator(256);
+            entry = new Entry(new BufferAllocator(256));
         } else {
-            bufferAllocator = FREE.removeFirst();
-            if (bufferAllocator.pointer == 0L) { // If the buffer was closed while in the pool
-                BUFFER_ALLOCATOR_ACCESS_TIME.removeLong(bufferAllocator);
-                bufferAllocator = new BufferAllocator(256);
+            entry = FREE.removeFirst();
+            if (entry.bufferAllocator.pointer == 0L) { // If the buffer was closed while in the pool
+                BUFFER_ALLOCATOR_MAPPING.remove(entry.bufferAllocator);
+                entry = new Entry(new BufferAllocator(256));
             }
         }
-        IN_USE.add(bufferAllocator);
-        BUFFER_ALLOCATOR_ACCESS_TIME.put(bufferAllocator, System.currentTimeMillis());
-        return bufferAllocator;
+        IN_USE.add(entry);
+        BUFFER_ALLOCATOR_MAPPING.put(entry.bufferAllocator, entry);
+        entry.onBorrow();
+        return entry.bufferAllocator;
     }
 
     public static void returnBufferAllocatorSafe(final BufferAllocator bufferAllocator) {
         RenderSystem.assertOnRenderThread();
-        if (!IN_USE.remove(bufferAllocator)) {
+        final Entry entry = BUFFER_ALLOCATOR_MAPPING.get(bufferAllocator);
+        if (!IN_USE.remove(entry)) {
             return;
         }
-        bufferAllocator.reset();
-        FREE.add(bufferAllocator);
+        entry.onReturn();
+        FREE.addFirst(entry);
     }
 
     public static int getSize() {
@@ -66,23 +68,59 @@ public static int getSize() {
 
     public static void onEndFrame() {
         if (!IN_USE.isEmpty()) {
-            ImmediatelyFast.LOGGER.warn(IN_USE.size() + " BufferAllocator(s) were not returned to the pool. Forcibly reclaiming them.");
-            for (BufferAllocator bufferAllocator : IN_USE) {
-                bufferAllocator.reset();
+            // Reclaim all buffers that were not returned to the pool this and the last frame
+            final boolean leak = IN_USE.removeIf(entry -> {
+                if (entry.inUseOverMultipleFrames) {
+                    entry.onReturn();
+                    FREE.addFirst(entry);
+                    return true;
+                }
+                return false;
+            });
+            if (leak) {
+                ImmediatelyFast.LOGGER.warn("Some BufferAllocators were not returned to the pool. Forcibly reclaiming them to prevent a memory leak.");
+            }
+
+            // Mark all as in use over multiple frames
+            for (Entry entry : IN_USE) {
+                entry.inUseOverMultipleFrames = true;
             }
-            FREE.addAll(IN_USE);
-            IN_USE.clear();
         }
-        BUFFER_ALLOCATOR_ACCESS_TIME.reference2LongEntrySet().removeIf(entry -> {
-            if (System.currentTimeMillis() - entry.getLongValue() > 60 * 1000) {
-                if (FREE.contains(entry.getKey())) {
-                    FREE.remove(entry.getKey());
-                    entry.getKey().close();
-                }
+
+        FREE.removeIf(entry -> {
+            if (entry.shouldBeClosed()) {
+                entry.bufferAllocator.close();
+                BUFFER_ALLOCATOR_MAPPING.remove(entry.bufferAllocator);
                 return true;
             }
             return false;
         });
     }
 
+    private static class Entry {
+
+        private final BufferAllocator bufferAllocator;
+        private long lastAccessTime;
+        private boolean inUseOverMultipleFrames;
+
+        public Entry(final BufferAllocator bufferAllocator) {
+            this.bufferAllocator = bufferAllocator;
+            this.lastAccessTime = System.currentTimeMillis();
+        }
+
+        public boolean shouldBeClosed() {
+            return System.currentTimeMillis() - this.lastAccessTime > 60 * 1000;
+        }
+
+        public void onBorrow() {
+            this.lastAccessTime = System.currentTimeMillis();
+        }
+
+        public void onReturn() {
+            this.bufferAllocator.reset();
+            this.inUseOverMultipleFrames = false;
+        }
+
+    }
+
 }