fix: memleak and locking issues

c2me-base/src/main/java/com/ishland/c2me/base/common/scheduler/ScheduledTask.java

@@ -2,7 +2,9 @@
 
 import com.ishland.flowsched.executor.LockToken;
 import com.ishland.flowsched.executor.Task;
+import it.unimi.dsi.fastutil.objects.ReferenceArrayList;
 
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
 
@@ -12,6 +14,7 @@ public class ScheduledTask<T> implements Task {
     private final Supplier<CompletableFuture<T>> action;
     private final LockToken[] lockTokens;
     private final CompletableFuture<T> future = new CompletableFuture<>();
+    private final ReferenceArrayList<Runnable> postExec = new ReferenceArrayList<>(4);
     private int priority = Integer.MAX_VALUE;
 
     public ScheduledTask(long pos, Supplier<CompletableFuture<T>> action, LockToken[] lockTokens) {
@@ -28,6 +31,13 @@ public void run() {
             } else {
                 future.complete(t);
             }
+            for (Runnable runnable : this.postExec) {
+                try {
+                    runnable.run();
+                } catch (Throwable t1) {
+                    t1.printStackTrace();
+                }
+            }
         });
     }
 
@@ -36,6 +46,12 @@ public void propagateException(Throwable t) {
         future.completeExceptionally(t);
     }
 
+    public void addPostExec(Runnable runnable) {
+        synchronized (this.postExec) {
+            postExec.add(Objects.requireNonNull(runnable));
+        }
+    }
+
     @Override
     public LockToken[] lockTokens() {
         return this.lockTokens;

c2me-base/src/main/java/com/ishland/c2me/base/common/scheduler/SchedulingManager.java

@@ -1,6 +1,7 @@
 package com.ishland.c2me.base.common.scheduler;
 
 import com.ishland.c2me.base.common.GlobalExecutors;
+import com.ishland.flowsched.structs.SimpleObjectPool;
 import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap;
 import it.unimi.dsi.fastutil.longs.Long2ReferenceOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectArraySet;
@@ -16,7 +17,9 @@ public class SchedulingManager {
 
     public static final int MAX_LEVEL = ChunkLevels.INACCESSIBLE + 1;
     private final Long2ReferenceOpenHashMap<ObjectArraySet<ScheduledTask<?>>> pos2Tasks = new Long2ReferenceOpenHashMap<>();
+    private final SimpleObjectPool<ObjectArraySet<ScheduledTask<?>>> pos2TasksPool = new SimpleObjectPool<>(unused -> new ObjectArraySet<>(), ObjectArraySet::clear, ObjectArraySet::clear, 2048);
     private final Long2IntOpenHashMap prioritiesFromLevel = new Long2IntOpenHashMap();
+    private final Object schedulingMutex = new Object();
     private final int id = COUNTER.getAndIncrement();
     private ChunkPos currentSyncLoad = null;
 
@@ -31,18 +34,30 @@ public SchedulingManager(Executor executor) {
     }
 
     public void enqueue(ScheduledTask<?> task) {
-        synchronized (this.prioritiesFromLevel) {
+        synchronized (this.schedulingMutex) {
             final long pos = task.getPos();
-            final ObjectArraySet<ScheduledTask<?>> locks = this.pos2Tasks.computeIfAbsent(pos, unused -> new ObjectArraySet<>());
+            final ObjectArraySet<ScheduledTask<?>> locks = this.pos2Tasks.computeIfAbsent(pos, unused -> this.pos2TasksPool.alloc());
             locks.add(task);
             updatePriorityInternal(pos);
         }
+        task.addPostExec(() -> {
+            synchronized (this.schedulingMutex) {
+                final ObjectArraySet<ScheduledTask<?>> tasks = this.pos2Tasks.get(task.getPos());
+                if (tasks != null) {
+                    tasks.remove(task);
+                    if (tasks.isEmpty()) {
+                        this.pos2Tasks.remove(task.getPos());
+                        this.pos2TasksPool.release(tasks);
+                    }
+                }
+            }
+        });
         GlobalExecutors.prioritizedScheduler.schedule(task);
     }
 
     public void updatePriorityFromLevel(long pos, int level) {
         this.executor.execute(() -> {
-            synchronized (this.prioritiesFromLevel) {
+            synchronized (this.schedulingMutex) {
                 if (prioritiesFromLevel.get(pos) == level) return;
                 if (level < MAX_LEVEL) {
                     prioritiesFromLevel.put(pos, level);
@@ -80,14 +95,16 @@ private void updatePriorityInternal(long pos) {
 
     public void setCurrentSyncLoad(ChunkPos pos) {
         executor.execute(() -> {
-            if (this.currentSyncLoad != null) {
-                final ChunkPos lastSyncLoad = this.currentSyncLoad;
-                this.currentSyncLoad = null;
-                updateSyncLoadInternal(lastSyncLoad);
-            }
-            if (pos != null) {
-                this.currentSyncLoad = pos;
-                updateSyncLoadInternal(pos);
+            synchronized (this.schedulingMutex) {
+                if (this.currentSyncLoad != null) {
+                    final ChunkPos lastSyncLoad = this.currentSyncLoad;
+                    this.currentSyncLoad = null;
+                    updateSyncLoadInternal(lastSyncLoad);
+                }
+                if (pos != null) {
+                    this.currentSyncLoad = pos;
+                    updateSyncLoadInternal(pos);
+                }
             }
         });
     }

c2me-notickvd/src/main/java/com/ishland/c2me/notickvd/common/PlayerNoTickDistanceMap.java

@@ -1,7 +1,7 @@
 package com.ishland.c2me.notickvd.common;
 
-import com.ishland.c2me.base.common.structs.DynamicPriorityQueue;
 import com.ishland.c2me.base.mixin.access.IThreadedAnvilChunkStorage;
+import com.ishland.flowsched.structs.DynamicPriorityQueue;
 import com.mojang.logging.LogUtils;
 import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap;
 import it.unimi.dsi.fastutil.longs.LongIterator;

c2me-opts-allocs/src/main/java/com/ishland/c2me/opts/allocs/common/PooledFeatureContext.java

@@ -1,6 +1,6 @@
 package com.ishland.c2me.opts.allocs.common;
 
-import com.ishland.c2me.base.common.structs.SimpleObjectPool;
+import com.ishland.flowsched.structs.SimpleObjectPool;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.random.Random;
 import net.minecraft.world.StructureWorldAccess;
@@ -13,7 +13,7 @@
 
 public class PooledFeatureContext<FC extends FeatureConfig> extends FeatureContext<FC> {
 
-    public static final ThreadLocal<SimpleObjectPool<PooledFeatureContext<?>>> POOL = ThreadLocal.withInitial(() -> new SimpleObjectPool<>(unused -> new PooledFeatureContext<>(), unused -> {}, 2048));
+    public static final ThreadLocal<SimpleObjectPool<PooledFeatureContext<?>>> POOL = ThreadLocal.withInitial(() -> new SimpleObjectPool<>(unused -> new PooledFeatureContext<>(), PooledFeatureContext::reInit, PooledFeatureContext::reInit, 2048));
 
     private Optional<ConfiguredFeature<?, ?>> feature;
     private StructureWorldAccess world;

c2me-opts-allocs/src/main/java/com/ishland/c2me/opts/allocs/mixin/object_pooling_caching/MixinConfiguredFeature.java

@@ -1,7 +1,7 @@
 package com.ishland.c2me.opts.allocs.mixin.object_pooling_caching;
 
 import com.ishland.c2me.opts.allocs.common.PooledFeatureContext;
-import com.ishland.c2me.base.common.structs.SimpleObjectPool;
+import com.ishland.flowsched.structs.SimpleObjectPool;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.random.Random;
 import net.minecraft.world.StructureWorldAccess;
@@ -36,7 +36,6 @@ public boolean generate(StructureWorldAccess world, ChunkGenerator chunkGenerato
             context.reInit(Optional.empty(), world, chunkGenerator, random, origin, this.config);
             return this.feature.generate(context);
         } finally {
-            context.reInit();
             pool.release(context);
         }
     }