Folia/patches/server/0003-Make-ChunkStatus.EMPTY-not-rely-on-the-main-thread-f.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Spottedleaf <Spottedleaf@users.noreply.github.com>
Date: Thu, 16 Feb 2023 16:50:05 -0800
Subject: [PATCH] Make ChunkStatus.EMPTY not rely on the main thread for
 completion

In order to do this, we need to push the POI consistency checks
to a later status. Since FULL is the only other status that
uses the main thread, it can go there.

The consistency checks are only really for when a desync occurs,
and so that delaying the check only matters when the chunk data
has desync'd. As long as the desync is sorted before the
chunk is full loaded (i.e before setBlock can occur on
a chunk), it should not matter.

This change is primarily due to behavioural changes
in the chunk task queue brought by region threading -
which is to split the queue into separate regions. As such,
it is required that in order for the sync load to complete
that the region owning the chunk drain and execute the task
while ticking. However, that is not always possible in
region threading. Thus, removing the main thread reliance allows
the chunk to progress without requiring a tick thread.
Specifically, this allows far sync loads (outside of a specific
regions bounds) to occur without issue - namely with structure
searching.

diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
index fb42d776f15f735fb59e972e00e2b512c23a8387..300700477ee34bc22b31315825c0e40f61070cd5 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkFullTask.java
@@ -2,6 +2,8 @@ package io.papermc.paper.chunk.system.scheduling;
 
 import ca.spottedleaf.concurrentutil.executor.standard.PrioritisedExecutor;
 import ca.spottedleaf.concurrentutil.util.ConcurrentUtil;
+import com.mojang.logging.LogUtils;
+import io.papermc.paper.chunk.system.poi.PoiChunk;
 import net.minecraft.server.level.ChunkMap;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.level.chunk.ChunkAccess;
@@ -9,10 +11,13 @@ import net.minecraft.world.level.chunk.ChunkStatus;
 import net.minecraft.world.level.chunk.ImposterProtoChunk;
 import net.minecraft.world.level.chunk.LevelChunk;
 import net.minecraft.world.level.chunk.ProtoChunk;
+import org.slf4j.Logger;
 import java.lang.invoke.VarHandle;
 
 public final class ChunkFullTask extends ChunkProgressionTask implements Runnable {
 
+    private static final Logger LOGGER = LogUtils.getClassLogger();
+
     protected final NewChunkHolder chunkHolder;
     protected final ChunkAccess fromChunk;
     protected final PrioritisedExecutor.PrioritisedTask convertToFullTask;
@@ -35,6 +40,15 @@ public final class ChunkFullTask extends ChunkProgressionTask implements Runnabl
         // See Vanilla protoChunkToFullChunk for what this function should be doing
         final LevelChunk chunk;
         try {
+            // moved from the load from nbt stage into here
+            final PoiChunk poiChunk = this.chunkHolder.getPoiChunk();
+            if (poiChunk == null) {
+                LOGGER.error("Expected poi chunk to be loaded with chunk for task " + this.toString());
+            } else {
+                poiChunk.load();
+                this.world.getPoiManager().checkConsistency(this.fromChunk);
+            }
+
             if (this.fromChunk instanceof ImposterProtoChunk wrappedFull) {
                 chunk = wrappedFull.getWrapped();
             } else {
diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkLoadTask.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkLoadTask.java
index be6f3f6a57668a9bd50d0ea5f2dd2335355b69d6..d52f68d0302d581b9bbe30fe0680fb8108e3da85 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkLoadTask.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkLoadTask.java
@@ -25,6 +25,7 @@ import org.slf4j.Logger;
 import java.lang.invoke.VarHandle;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 
 public final class ChunkLoadTask extends ChunkProgressionTask {
@@ -34,9 +35,11 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
     private final NewChunkHolder chunkHolder;
     private final ChunkDataLoadTask loadTask;
 
-    private boolean cancelled;
+    private volatile boolean cancelled;
     private NewChunkHolder.GenericDataLoadTaskCallback entityLoadTask;
     private NewChunkHolder.GenericDataLoadTaskCallback poiLoadTask;
+    private GenericDataLoadTask.TaskResult<ChunkAccess, Throwable> loadResult;
+    private final AtomicInteger taskCountToComplete = new AtomicInteger(3); // one for poi, one for entity, and one for chunk data
 
     protected ChunkLoadTask(final ChunkTaskScheduler scheduler, final ServerLevel world, final int chunkX, final int chunkZ,
                             final NewChunkHolder chunkHolder, final PrioritisedExecutor.Priority priority) {
@@ -44,10 +47,18 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
         this.chunkHolder = chunkHolder;
         this.loadTask = new ChunkDataLoadTask(scheduler, world, chunkX, chunkZ, priority);
         this.loadTask.addCallback((final GenericDataLoadTask.TaskResult<ChunkAccess, Throwable> result) -> {
-            ChunkLoadTask.this.complete(result == null ? null : result.left(), result == null ? null : result.right());
+            ChunkLoadTask.this.loadResult = result; // must be before getAndDecrement
+            ChunkLoadTask.this.tryCompleteLoad();
         });
     }
 
+    private void tryCompleteLoad() {
+        if (this.taskCountToComplete.decrementAndGet() == 0) {
+            final GenericDataLoadTask.TaskResult<ChunkAccess, Throwable> result = this.cancelled ? null : this.loadResult; // only after the getAndDecrement
+            ChunkLoadTask.this.complete(result == null ? null : result.left(), result == null ? null : result.right());
+        }
+    }
+
     @Override
     public ChunkStatus getTargetStatus() {
         return ChunkStatus.EMPTY;
@@ -65,11 +76,8 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
         final NewChunkHolder.GenericDataLoadTaskCallback entityLoadTask;
         final NewChunkHolder.GenericDataLoadTaskCallback poiLoadTask;
 
-        final AtomicInteger count = new AtomicInteger();
         final Consumer<GenericDataLoadTask.TaskResult<?, ?>> scheduleLoadTask = (final GenericDataLoadTask.TaskResult<?, ?> result) -> {
-            if (count.decrementAndGet() == 0) {
-                ChunkLoadTask.this.loadTask.schedule(false);
-            }
+            ChunkLoadTask.this.tryCompleteLoad();
         };
 
         // NOTE: it is IMPOSSIBLE for getOrLoadEntityData/getOrLoadPoiData to complete synchronously, because
@@ -85,16 +93,16 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
             }
             if (!this.chunkHolder.isEntityChunkNBTLoaded()) {
                 entityLoadTask = this.chunkHolder.getOrLoadEntityData((Consumer)scheduleLoadTask);
-                count.setPlain(count.getPlain() + 1);
             } else {
                 entityLoadTask = null;
+                this.taskCountToComplete.getAndDecrement(); // we know the chunk load is not done here, as it is not scheduled
             }
 
             if (!this.chunkHolder.isPoiChunkLoaded()) {
                 poiLoadTask = this.chunkHolder.getOrLoadPoiData((Consumer)scheduleLoadTask);
-                count.setPlain(count.getPlain() + 1);
             } else {
                 poiLoadTask = null;
+                this.taskCountToComplete.getAndDecrement(); // we know the chunk load is not done here, as it is not scheduled
             }
 
             this.entityLoadTask = entityLoadTask;
@@ -107,14 +115,11 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
             entityLoadTask.schedule();
         }
 
-        if (poiLoadTask !=  null) {
+        if (poiLoadTask != null) {
             poiLoadTask.schedule();
         }
 
-        if (entityLoadTask == null && poiLoadTask == null) {
-            // no need to wait on those, we can schedule now
-            this.loadTask.schedule(false);
-        }
+        this.loadTask.schedule(false);
     }
 
     @Override
@@ -129,15 +134,20 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
 
         /*
         Note: The entityLoadTask/poiLoadTask do not complete when cancelled,
-        but this is fine because if they are successfully cancelled then
-        we will successfully cancel the load task, which will complete when cancelled
+        so we need to manually try to complete in those cases
+        It is also important to note that we set the cancelled field first, just in case
+        the chunk load task attempts to complete with a non-null value
         */
 
         if (this.entityLoadTask != null) {
-            this.entityLoadTask.cancel();
+            if (this.entityLoadTask.cancel()) {
+                this.tryCompleteLoad();
+            }
         }
         if (this.poiLoadTask != null) {
-            this.poiLoadTask.cancel();
+            if (this.poiLoadTask.cancel()) {
+                this.tryCompleteLoad();
+            }
         }
         this.loadTask.cancel();
     }
@@ -249,7 +259,7 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
         }
     }
 
-    public final class ChunkDataLoadTask extends CallbackDataLoadTask<ChunkSerializer.InProgressChunkHolder, ChunkAccess> {
+    public static final class ChunkDataLoadTask extends CallbackDataLoadTask<ChunkAccess, ChunkAccess> {
         protected ChunkDataLoadTask(final ChunkTaskScheduler scheduler, final ServerLevel world, final int chunkX,
                                     final int chunkZ, final PrioritisedExecutor.Priority priority) {
             super(scheduler, world, chunkX, chunkZ, RegionFileIOThread.RegionFileType.CHUNK_DATA, priority);
@@ -262,7 +272,7 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
 
         @Override
         protected boolean hasOnMain() {
-            return true;
+            return false;
         }
 
         @Override
@@ -272,35 +282,30 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
 
         @Override
         protected PrioritisedExecutor.PrioritisedTask createOnMain(final Runnable run, final PrioritisedExecutor.Priority priority) {
-            return this.scheduler.createChunkTask(this.chunkX, this.chunkZ, run, priority);
+            throw new UnsupportedOperationException();
         }
 
         @Override
-        protected TaskResult<ChunkAccess, Throwable> completeOnMainOffMain(final ChunkSerializer.InProgressChunkHolder data, final Throwable throwable) {
-            if (data != null) {
-                return null;
-            }
-
-            final PoiChunk poiChunk = ChunkLoadTask.this.chunkHolder.getPoiChunk();
-            if (poiChunk == null) {
-                LOGGER.error("Expected poi chunk to be loaded with chunk for task " + this.toString());
-            } else if (!poiChunk.isLoaded()) {
-                // need to call poiChunk.load() on main
-                return null;
-            }
+        protected TaskResult<ChunkAccess, Throwable> completeOnMainOffMain(final ChunkAccess data, final Throwable throwable) {
+            throw new UnsupportedOperationException();
+        }
 
-            return new TaskResult<>(this.getEmptyChunk(), null);
+        private ProtoChunk getEmptyChunk() {
+            return new ProtoChunk(
+                new ChunkPos(this.chunkX, this.chunkZ), UpgradeData.EMPTY, this.world,
+                this.world.registryAccess().registryOrThrow(Registries.BIOME), (BlendingData)null
+            );
         }
 
         @Override
-        protected TaskResult<ChunkSerializer.InProgressChunkHolder, Throwable> runOffMain(final CompoundTag data, final Throwable throwable) {
+        protected TaskResult<ChunkAccess, Throwable> runOffMain(final CompoundTag data, final Throwable throwable) {
             if (throwable != null) {
                 LOGGER.error("Failed to load chunk data for task: " + this.toString() + ", chunk data will be lost", throwable);
-                return new TaskResult<>(null, null);
+                return new TaskResult<>(this.getEmptyChunk(), null);
             }
 
             if (data == null) {
-                return new TaskResult<>(null, null);
+                return new TaskResult<>(this.getEmptyChunk(), null);
             }
 
             // need to convert data, and then deserialize it
@@ -319,53 +324,18 @@ public final class ChunkLoadTask extends ChunkProgressionTask {
                     this.world, chunkMap.getPoiManager(), chunkPos, converted, true
                 );
 
-                return new TaskResult<>(chunkHolder, null);
+                return new TaskResult<>(chunkHolder.protoChunk, null);
             } catch (final ThreadDeath death) {
                 throw death;
             } catch (final Throwable thr2) {
                 LOGGER.error("Failed to parse chunk data for task: " + this.toString() + ", chunk data will be lost", thr2);
-                return new TaskResult<>(null, thr2);
+                return new TaskResult<>(this.getEmptyChunk(), thr2);
             }
         }
 
-        private ProtoChunk getEmptyChunk() {
-            return new ProtoChunk(
-                new ChunkPos(this.chunkX, this.chunkZ), UpgradeData.EMPTY, this.world,
-                this.world.registryAccess().registryOrThrow(Registries.BIOME), (BlendingData)null
-            );
-        }
-
         @Override
-        protected TaskResult<ChunkAccess, Throwable> runOnMain(final ChunkSerializer.InProgressChunkHolder data, final Throwable throwable) {
-            final PoiChunk poiChunk = ChunkLoadTask.this.chunkHolder.getPoiChunk();
-            if (poiChunk == null) {
-                LOGGER.error("Expected poi chunk to be loaded with chunk for task " + this.toString());
-            } else {
-                poiChunk.load();
-            }
-
-            if (data == null || data.protoChunk == null) {
-                // throwable could be non-null, but the off-main task will print its exceptions - so we don't need to care,
-                // it's handled already
-
-                return new TaskResult<>(this.getEmptyChunk(), null);
-            }
-
-            // have tasks to run (at this point, it's just the POI consistency checking)
-            try {
-                if (data.tasks != null) {
-                    for (int i = 0, len = data.tasks.size(); i < len; ++i) {
-                        data.tasks.poll().run();
-                    }
-                }
-
-                return new TaskResult<>(data.protoChunk, null);
-            } catch (final ThreadDeath death) {
-                throw death;
-            } catch (final Throwable thr2) {
-                LOGGER.error("Failed to parse main tasks for task " + this.toString() + ", chunk data will be lost", thr2);
-                return new TaskResult<>(this.getEmptyChunk(), null);
-            }
+        protected TaskResult<ChunkAccess, Throwable> runOnMain(final ChunkAccess data, final Throwable throwable) {
+            throw new UnsupportedOperationException();
         }
     }
 
diff --git a/src/main/java/net/minecraft/world/entity/ai/village/poi/PoiManager.java b/src/main/java/net/minecraft/world/entity/ai/village/poi/PoiManager.java
index 8950b220b9a3512cd4667beb7bdec0e82e07edc6..9be85eb0abec02bc0e0eded71c34ab1c565c63e7 100644
--- a/src/main/java/net/minecraft/world/entity/ai/village/poi/PoiManager.java
+++ b/src/main/java/net/minecraft/world/entity/ai/village/poi/PoiManager.java
@@ -328,6 +328,12 @@ public class PoiManager extends SectionStorage<PoiSection> {
             }
         }
     }
+
+    public void checkConsistency(net.minecraft.world.level.chunk.ChunkAccess chunk) {
+        for (LevelChunkSection section : chunk.getSections()) {
+            this.checkConsistencyWithBlocks(chunk.getPos(), section);
+        }
+    }
     // Paper end - rewrite chunk system
 
     public void checkConsistencyWithBlocks(ChunkPos chunkPos, LevelChunkSection chunkSection) {
diff --git a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
index 0ec80b83a99bfdb1f985045d98a81905a8a5a3ac..9d6f4749ed72fe319754ccea28f20fa97286527d 100644
--- a/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
+++ b/src/main/java/net/minecraft/world/level/chunk/storage/ChunkSerializer.java
@@ -122,13 +122,11 @@ public class ChunkSerializer {
     public static final class InProgressChunkHolder {
 
         public final ProtoChunk protoChunk;
-        public final java.util.ArrayDeque<Runnable> tasks;
 
         public CompoundTag poiData;
 
-        public InProgressChunkHolder(final ProtoChunk protoChunk, final java.util.ArrayDeque<Runnable> tasks) {
+        public InProgressChunkHolder(final ProtoChunk protoChunk) {
             this.protoChunk = protoChunk;
-            this.tasks = tasks;
         }
     }
     // Paper end
@@ -136,7 +134,6 @@ public class ChunkSerializer {
     public static ProtoChunk read(ServerLevel world, PoiManager poiStorage, ChunkPos chunkPos, CompoundTag nbt) {
         // Paper start - add variant for async calls
         InProgressChunkHolder holder = loadChunk(world, poiStorage, chunkPos, nbt, true);
-        holder.tasks.forEach(Runnable::run);
         return holder.protoChunk;
     }
 
@@ -145,7 +142,6 @@ public class ChunkSerializer {
     private static final boolean JUST_CORRUPT_IT = Boolean.getBoolean("Paper.ignoreWorldDataVersion");
     // Paper end
     public static InProgressChunkHolder loadChunk(ServerLevel world, PoiManager poiStorage, ChunkPos chunkPos, CompoundTag nbt, boolean distinguish) {
-        java.util.ArrayDeque<Runnable> tasksToExecuteOnMain = new java.util.ArrayDeque<>();
         // Paper end
         // Paper start - Do NOT attempt to load chunks saved with newer versions
         if (nbt.contains("DataVersion", 99)) {
@@ -223,9 +219,7 @@ public class ChunkSerializer {
                 LevelChunkSection chunksection = new LevelChunkSection(b0, datapaletteblock, (PalettedContainer) object); // CraftBukkit - read/write
 
                 achunksection[k] = chunksection;
-                tasksToExecuteOnMain.add(() -> { // Paper - delay this task since we're executing off-main
-                poiStorage.checkConsistencyWithBlocks(chunkPos, chunksection);
-                }); // Paper - delay this task since we're executing off-main
+                // Paper - rewrite chunk system - moved to final load stage
             }
 
             boolean flag3 = nbttagcompound1.contains("BlockLight", 7);
@@ -403,7 +397,7 @@ public class ChunkSerializer {
         }
 
         if (chunkstatus_type == ChunkStatus.ChunkType.LEVELCHUNK) {
-            return new InProgressChunkHolder(new ImposterProtoChunk((LevelChunk) object1, false), tasksToExecuteOnMain); // Paper - Async chunk loading
+            return new InProgressChunkHolder(new ImposterProtoChunk((LevelChunk) object1, false)); // Paper - Async chunk loading
         } else {
             ProtoChunk protochunk1 = (ProtoChunk) object1;
 
@@ -446,7 +440,7 @@ public class ChunkSerializer {
                 protochunk1.setCarvingMask(worldgenstage_features, new CarvingMask(nbttagcompound4.getLongArray(s1), ((ChunkAccess) object1).getMinBuildHeight()));
             }
 
-            return new InProgressChunkHolder(protochunk1, tasksToExecuteOnMain); // Paper - Async chunk loading
+            return new InProgressChunkHolder(protochunk1); // Paper - Async chunk loading
         }
     }