fix(scans): limits not persisting across restarts

Author: Jsinco

Insights-API/src/main/java/dev/frankheijden/insights/api/concurrent/ChunkContainerExecutor.java

@@ -86,7 +86,7 @@ public CompletableFuture<Storage> submit(ChunkContainer container, ScanOptions o
             scanTracker.set(worldUid, chunkKey, true);
         }
 
-        return submit(container).thenApply(storage -> {
+        return container.scan(containerExecutor).thenApply(storage -> {
             if (options.save()) worldStorage.getWorld(worldUid).put(chunkKey, storage);
             if (options.track()) scanTracker.set(worldUid, chunkKey, false);
 
@@ -108,6 +108,11 @@ public CompletableFuture<Void> submit(RunnableContainer container) {
         return containerExecutor.submit(container);
     }
 
+    @Override
+    public int getTimeout() {
+        return containerExecutor.getTimeout();
+    }
+
     @Override
     public void shutdown() {
         containerExecutor.shutdown();

Insights-API/src/main/java/dev/frankheijden/insights/api/concurrent/ContainerExecutor.java

@@ -10,5 +10,11 @@ public interface ContainerExecutor {
 
     CompletableFuture<Void> submit(RunnableContainer container);
 
+    /**
+     * Timeout in milliseconds.
+     * @return Timeout in milliseconds.
+     */
+    int getTimeout();
+
     void shutdown();
 }

Insights-API/src/main/java/dev/frankheijden/insights/api/concurrent/containers/ChunkContainer.java

@@ -1,5 +1,6 @@
 package dev.frankheijden.insights.api.concurrent.containers;
 
+import dev.frankheijden.insights.api.concurrent.ContainerExecutor;
 import dev.frankheijden.insights.api.concurrent.ScanOptions;
 import dev.frankheijden.insights.api.concurrent.storage.DistributionStorage;
 import dev.frankheijden.insights.api.exceptions.ChunkIOException;
@@ -16,6 +17,7 @@
 import java.io.IOException;
 import java.util.EnumMap;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 
 public abstract class ChunkContainer implements SupplierContainer<DistributionStorage> {
@@ -74,6 +76,16 @@ public ChunkCuboid getChunkCuboid() {
 
     public abstract void getChunkEntities(Consumer<@NotNull ChunkEntity> entityConsumer) throws IOException;
 
+    /**
+     * Implementations may vary what thread this method performs operations on.
+     *
+     * @see LoadedChunkContainer#scan(ContainerExecutor)
+     * @see UnloadedChunkContainer#scan(ContainerExecutor)
+     * @param executor The executor to run the scan on.
+     * @return A CompletableFuture of the DistributionStorage.
+     */
+    public abstract CompletableFuture<DistributionStorage> scan(ContainerExecutor executor);
+
     @Override
     public DistributionStorage get() {
         ChunkVector min = cuboid.getMin();

Insights-API/src/main/java/dev/frankheijden/insights/api/concurrent/containers/LoadedChunkContainer.java

@@ -1,7 +1,9 @@
 package dev.frankheijden.insights.api.concurrent.containers;
 
 import dev.frankheijden.insights.api.InsightsPlugin;
+import dev.frankheijden.insights.api.concurrent.ContainerExecutor;
 import dev.frankheijden.insights.api.concurrent.ScanOptions;
+import dev.frankheijden.insights.api.concurrent.storage.DistributionStorage;
 import dev.frankheijden.insights.api.objects.chunk.ChunkCuboid;
 import dev.frankheijden.insights.nms.core.ChunkEntity;
 import dev.frankheijden.insights.nms.core.ChunkSection;
@@ -11,6 +13,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
 public class LoadedChunkContainer extends ChunkContainer {
@@ -27,39 +30,28 @@ public LoadedChunkContainer(InsightsNMS nms, Chunk chunk, ChunkCuboid cuboid, Sc
 
     @Override
     public void getChunkSections(Consumer<@Nullable ChunkSection> sectionConsumer) {
-        runOnOwningRegion(() -> nms.getLoadedChunkSections(chunk, sectionConsumer));
+        nms.getLoadedChunkSections(chunk, sectionConsumer);
     }
 
     @Override
     public void getChunkEntities(Consumer<@NotNull ChunkEntity> entityConsumer) {
-        runOnOwningRegion(() -> nms.getLoadedChunkEntities(chunk, entityConsumer));
+        nms.getLoadedChunkEntities(chunk, entityConsumer);
     }
 
-    /**
-     * Runs the given chunk-reading task on the region thread which owns this chunk, blocking the
-     * calling (worker) thread until it has completed.
-     *
-     * <p>This synchronization is required: {@link ChunkContainer#get()} reads the populated
-     * material/entity maps immediately after invoking these methods. If the read were merely
-     * scheduled on the region thread (fire-and-forget), {@code get()} would observe still-empty
-     * maps and return a distribution counting 0, causing limits to be bypassed on every chunk
-     * that is rescanned after being (re)loaded.
-     */
-    private void runOnOwningRegion(Runnable task) {
+    @Override
+    public CompletableFuture<DistributionStorage> scan(ContainerExecutor executor) {
         if (Bukkit.isOwnedByCurrentRegion(chunk.getWorld(), chunkX, chunkZ)) {
-            task.run();
-            return;
+            return CompletableFuture.completedFuture(this.get());
         }
 
-        CompletableFuture<Void> future = new CompletableFuture<>();
+        CompletableFuture<DistributionStorage> future = new CompletableFuture<>();
         Bukkit.getRegionScheduler().execute(InsightsPlugin.getInstance(), chunk.getWorld(), chunkX, chunkZ, () -> {
             try {
-                task.run();
-                future.complete(null);
+                future.complete(this.get());
             } catch (Throwable th) {
                 future.completeExceptionally(th);
             }
         });
-        future.join();
+        return future.orTimeout(executor.getTimeout(), TimeUnit.MILLISECONDS);
     }
 }

Insights-API/src/main/java/dev/frankheijden/insights/api/concurrent/containers/UnloadedChunkContainer.java

@@ -1,6 +1,8 @@
 package dev.frankheijden.insights.api.concurrent.containers;
 
+import dev.frankheijden.insights.api.concurrent.ContainerExecutor;
 import dev.frankheijden.insights.api.concurrent.ScanOptions;
+import dev.frankheijden.insights.api.concurrent.storage.DistributionStorage;
 import dev.frankheijden.insights.api.objects.chunk.ChunkCuboid;
 import dev.frankheijden.insights.nms.core.ChunkEntity;
 import dev.frankheijden.insights.nms.core.ChunkSection;
@@ -9,6 +11,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 
 public class UnloadedChunkContainer extends ChunkContainer {
@@ -36,4 +39,9 @@ public void getChunkSections(Consumer<@Nullable ChunkSection> sectionConsumer) {
     public void getChunkEntities(Consumer<@NotNull ChunkEntity> entityConsumer) throws IOException {
         nms.getUnloadedChunkEntities(world, chunkX, chunkZ, entityConsumer);
     }
+
+    @Override
+    public CompletableFuture<DistributionStorage> scan(ContainerExecutor executor) {
+        return executor.submit(this);
+    }
 }

Insights-Core/src/main/java/dev/frankheijden/insights/concurrent/ContainerExecutorService.java

@@ -54,6 +54,11 @@ public CompletableFuture<Void> submit(RunnableContainer container) {
         return CompletableFuture.runAsync(container, executor).orTimeout(timeoutMs, TimeUnit.MILLISECONDS);
     }
 
+    @Override
+    public int getTimeout() {
+        return timeoutMs;
+    }
+
     public int getQueueSize() {
         return executor.getQueue().size();
     }