Solving the concurrent object heap issue (and destruction from other

threads issue)
MCRcortex · May 22, 2024 · 2e98856 · 2e98856
1 parent d9d1b13
commit 2e98856
Show file tree

Hide file tree

Showing 9 changed files with 70 additions and 30 deletions.
diff --git a/src/main/java/me/cortex/vulkanite/acceleration/AccelerationBlasBuilder.java b/src/main/java/me/cortex/vulkanite/acceleration/AccelerationBlasBuilder.java
@@ -8,6 +8,7 @@
 import me.cortex.vulkanite.compat.IAccelerationBuildResult;
 import me.cortex.vulkanite.lib.base.VContext;
 import me.cortex.vulkanite.lib.base.VRef;
+import me.cortex.vulkanite.lib.base.VRegistry;
 import me.cortex.vulkanite.lib.cmd.VCmdBuff;
 import me.cortex.vulkanite.lib.cmd.VCommandPool;
 import me.cortex.vulkanite.lib.descriptors.DescriptorSetLayoutBuilder;
@@ -171,6 +172,8 @@ private void run() {
                 }
             }
 
+            VRegistry.INSTANCE.threadLocalCollect();
+
             var sinlgeUsePoolWorker = context.cmd.getSingleUsePool();
 
             //Jobs are batched and built on the async vulkan queue then block synchronized with fence

diff --git a/src/main/java/me/cortex/vulkanite/acceleration/AccelerationManager.java b/src/main/java/me/cortex/vulkanite/acceleration/AccelerationManager.java
@@ -60,8 +60,7 @@ public void updateTick() {
     public VRef<VAccelerationStructure> buildTLAS(int queueId, VCmdBuff cmd) {
         blasExecutions.forEach(exec -> ctx.cmd.queueWaitForExeuction(queueId, blasBuilder.getAsyncQueue(), exec));
         blasExecutions.clear();
-        tlasManager.buildTLAS(cmd);
-        return tlasManager.getTlas();
+        return tlasManager.buildTLAS(cmd);
     }
 
     public void sectionRemove(RenderSection section) {

diff --git a/src/main/java/me/cortex/vulkanite/acceleration/AccelerationTLASManager.java b/src/main/java/me/cortex/vulkanite/acceleration/AccelerationTLASManager.java
@@ -38,8 +38,6 @@ public class AccelerationTLASManager {
     private final VContext context;
     private final int queue;
 
-    private VRef<VAccelerationStructure> currentTLAS;
-
     public AccelerationTLASManager(VContext context, int queue) {
         this.context = context;
         this.queue = queue;
@@ -71,7 +69,7 @@ public void removeSection(RenderSection section) {
 
     // TODO: cleanup, this is very messy
     // FIXME: in the case of no geometry create an empty tlas or something???
-    public void buildTLAS(VCmdBuff cmd) {
+    public VRef<VAccelerationStructure> buildTLAS(VCmdBuff cmd) {
         RenderSystem.assertOnRenderThread();
 
         // NOTE: renderLink is required to ensure that we are not overriding memory that
@@ -188,17 +186,10 @@ public void buildTLAS(VCmdBuff cmd) {
 
             cmd.encodeMemoryBarrier();
 
-            if (currentTLAS != null) {
-                currentTLAS.close();
-            }
-            currentTLAS = tlas;
+            return tlas;
         }
     }
 
-    public VRef<VAccelerationStructure> getTlas() {
-        return currentTLAS.addRef();
-    }
-
 
     // Manages entries in the VkAccelerationStructureInstanceKHR buffer, ment to
     // reuse as much as possible and be very efficient
@@ -373,6 +364,7 @@ public void resizeBindlessSet(int newSize) {
                                 .descriptorCount(setCapacity);
                         vkUpdateDescriptorSets(context.device, null, setCopy);
                     }
+                    geometryBufferDescSet.close();
                 }
 
                 geometryBufferDescSet = newGeometryBufferDescSet;

diff --git a/src/main/java/me/cortex/vulkanite/client/rendering/VulkanPipeline.java b/src/main/java/me/cortex/vulkanite/client/rendering/VulkanPipeline.java
@@ -243,6 +243,7 @@ private void buildEntities() {
     public void renderPostShadows(List<VRef<VGImage>> vgOutImgs, Camera camera, ShaderStorageBuffer[] ssbos, MixinCelestialUniforms celestialUniforms) {
         ctx.cmd.newFrame();
 
+        VRegistry.INSTANCE.threadLocalCollect();
         System.out.println(VRegistry.INSTANCE.dumpStats());
 
         buildEntities();

diff --git a/src/main/java/me/cortex/vulkanite/lib/base/VObject.java b/src/main/java/me/cortex/vulkanite/lib/base/VObject.java
@@ -16,8 +16,9 @@ protected void incRef() {
 
     protected void decRef() {
         if (refCount.decrementAndGet() == 0) {
-            free();
             VRegistry.INSTANCE.unregister(this);
         }
     }
+
+    public Object heap = null;
 }
diff --git a/src/main/java/me/cortex/vulkanite/lib/base/VRegistry.java b/src/main/java/me/cortex/vulkanite/lib/base/VRegistry.java
@@ -1,11 +1,10 @@
 package me.cortex.vulkanite.lib.base;
 
 import com.google.common.collect.ConcurrentHashMultiset;
+import java.util.concurrent.ConcurrentLinkedDeque;
 import me.cortex.vulkanite.lib.memory.VBuffer;
 
-import java.util.HashMap;
-import java.util.Map;
-import java.util.TreeMap;
+import java.util.*;
 
 import static org.lwjgl.vulkan.KHRAccelerationStructure.VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR;
 import static org.lwjgl.vulkan.KHRAccelerationStructure.VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR;
@@ -16,17 +15,44 @@
 public class VRegistry {
     public static final VRegistry INSTANCE = new VRegistry();
 
-    private final ConcurrentHashMultiset<VObject> objects = ConcurrentHashMultiset.create();
+    private static class ObjectHeap {
+        protected final long threadId = Thread.currentThread().getId();
+        protected final String threadName = Thread.currentThread().getName();
+
+        protected final HashSet<VObject> objects = new HashSet<>();
+        protected final ConcurrentLinkedDeque<VObject> toFree = new ConcurrentLinkedDeque<>();
+
+        protected void collect() {
+            if (Thread.currentThread().getId() != threadId) {
+                throw new IllegalStateException("Object heap collect called from wrong thread");
+            }
+            VObject object;
+            while ((object = toFree.poll()) != null) {
+                objects.remove(object);
+                object.free();
+            }
+        }
+    }
+
+    private final ConcurrentHashMultiset<ObjectHeap> objectHeaps = ConcurrentHashMultiset.create();
+
+    private final ThreadLocal<ObjectHeap> heap = ThreadLocal.withInitial(()->{
+        var heap = new ObjectHeap();
+        objectHeaps.add(heap);
+        return heap;
+    });
 
     private VRegistry() {
     }
 
     public void register(VObject object) {
-        objects.add(object);
+        heap.get().objects.add(object);
+        object.heap = heap.get();
     }
 
     public void unregister(VObject object) {
-        objects.remove(object);
+        ObjectHeap heap = (ObjectHeap) object.heap;
+        heap.toFree.add(object);
     }
 
     private static final HashMap<Integer, String> usageNames = new HashMap<>() {{
@@ -45,7 +71,11 @@ public void unregister(VObject object) {
     public String dumpStats() {
         final StringBuilder sb = new StringBuilder();
 
-        sb.append("VRegistry: ").append(objects.size()).append(" objects\n");
+        ObjectHeap heap = this.heap.get();
+        var objects = heap.objects;
+
+        sb.append("\nVRegistry (Thread=").append(heap.threadId).append(" ").append(heap.threadName).append("): ");
+        sb.append(objects.size()).append(" objects\n");
         sb.append("Objects:\n");
 
         Map<String, Integer> typeCount = new TreeMap<>();
@@ -80,4 +110,8 @@ public String dumpStats() {
 
         return sb.toString();
     }
+
+    public void threadLocalCollect() {
+        heap.get().collect();
+    }
 }
diff --git a/src/main/java/me/cortex/vulkanite/lib/descriptors/DescriptorUpdateBuilder.java b/src/main/java/me/cortex/vulkanite/lib/descriptors/DescriptorUpdateBuilder.java
@@ -76,7 +76,7 @@ public DescriptorUpdateBuilder buffer(int binding, final VRef<VBuffer> buffer, l
         if (refSet != null && refSet.getBindingAt(binding) == null) {
             return this;
         }
-        setRef.get().refs.put(binding, buffer.addRefGeneric());
+        setRef.get().addRef(binding, buffer.addRefGeneric());
         updates.get()
                 .sType$Default()
                 .dstBinding(binding)
@@ -98,7 +98,7 @@ public DescriptorUpdateBuilder buffer(int binding, int dstArrayElement, final Li
         }
         var bufInfo = VkDescriptorBufferInfo.calloc(buffers.size());
         for (int i = 0; i < buffers.size(); i++) {
-            setRef.get().refs.put(binding, buffers.get(i).addRefGeneric());
+            setRef.get().addRef(binding, buffers.get(i).addRefGeneric());
             bufInfo.get(i)
                     .buffer(buffers.get(i).get().buffer())
                     .offset(0)
@@ -124,7 +124,7 @@ public DescriptorUpdateBuilder uniform(int binding, final VRef<VBuffer> buffer,
         if (refSet != null && refSet.getBindingAt(binding) == null) {
             return this;
         }
-        setRef.get().refs.put(binding, buffer.addRefGeneric());
+        setRef.get().addRef(binding, buffer.addRefGeneric());
         updates.get()
                 .sType$Default()
                 .dstBinding(binding)
@@ -145,7 +145,7 @@ public DescriptorUpdateBuilder acceleration(int binding, VRef<VAccelerationStruc
         }
         var buff = stack.mallocLong(structures.length);
         for (var structure : structures) {
-            setRef.get().refs.put(binding, structure.addRefGeneric());
+            setRef.get().addRef(binding, structure.addRefGeneric());
             buff.put(structure.get().structure);
         }
         buff.rewind();
@@ -167,7 +167,7 @@ public DescriptorUpdateBuilder imageStore(int binding, int dstArrayElement, fina
         }
         var imgInfo = VkDescriptorImageInfo.calloc(views.size());
         for (int i = 0; i < views.size(); i++) {
-            setRef.get().refs.put(binding, views.get(i).addRefGeneric());
+            setRef.get().addRef(binding, views.get(i).addRefGeneric());
             imgInfo.get(i)
                     .imageLayout(VK_IMAGE_LAYOUT_GENERAL)
                     .imageView(viewOrPlaceholder(views.get(i)));
@@ -189,7 +189,7 @@ public DescriptorUpdateBuilder imageStore(int binding, int layout, final VRef<VI
         if (refSet != null && refSet.getBindingAt(binding) == null) {
             return this;
         }
-        setRef.get().refs.put(binding, view.addRefGeneric());
+        setRef.get().addRef(binding, view.addRefGeneric());
         updates.get()
                 .sType$Default()
                 .dstBinding(binding)
@@ -211,7 +211,7 @@ public DescriptorUpdateBuilder imageSampler(int binding, int layout, final VRef<
         if (refSet != null && refSet.getBindingAt(binding) == null) {
             return this;
         }
-        setRef.get().refs.put(binding, view.addRefGeneric());
+        setRef.get().addRef(binding, view.addRefGeneric());
         updates.get()
                 .sType$Default()
                 .dstBinding(binding)

diff --git a/src/main/java/me/cortex/vulkanite/lib/descriptors/VDescriptorSet.java b/src/main/java/me/cortex/vulkanite/lib/descriptors/VDescriptorSet.java
@@ -11,7 +11,14 @@ public class VDescriptorSet extends VObject {
     public final long poolHandle;
     public final long set;
 
-    public Map<Integer, VRef<VObject>> refs = new HashMap<>();
+    private Map<Integer, VRef<VObject>> refs = new HashMap<>();
+
+    public void addRef(int binding, VRef<VObject> ref) {
+        var old = refs.put(binding, ref);
+        if (old != null) {
+            old.close();
+        }
+    }
 
     protected VDescriptorSet(VRef<VDescriptorPool> pool, long poolHandle, long set) {
         this.pool = pool;
@@ -22,5 +29,6 @@ protected VDescriptorSet(VRef<VDescriptorPool> pool, long poolHandle, long set)
     @Override
     protected void free() {
         pool.get().freeSet(this);
+        refs.values().forEach(VRef::close);
     }
 }
diff --git a/src/main/java/me/cortex/vulkanite/lib/memory/MemoryManager.java b/src/main/java/me/cortex/vulkanite/lib/memory/MemoryManager.java
@@ -303,7 +303,9 @@ public VRef<VAccelerationStructure> createAcceleration(long size, int alignment,
                             .size(size)
                             .buffer(buffer.get().buffer()), null, pAccelerationStructure),
                     "Failed to create acceleration acceleration structure");
-            return VAccelerationStructure.create(device, pAccelerationStructure.get(0), buffer);
+            var ret = VAccelerationStructure.create(device, pAccelerationStructure.get(0), buffer);
+            buffer.close();
+            return ret;
         }
     }