Address POI memory leak and Generate Command issues (#937)

Co-authored-by: embeddedt <[email protected]>

patches/net/minecraft/server/level/ChunkMap.java.patch

@@ -8,8 +8,12 @@
              return p_140179_;
          }
      }
-@@ -535,6 +_,7 @@
+@@ -533,8 +_,11 @@
+                 this.scheduleUnload(p_140182_, p_140183_);
+             } else {
                  if (this.pendingUnloads.remove(p_140182_, p_140183_) && p_203002_ != null) {
++                    net.neoforged.neoforge.common.CommonHooks.onChunkUnload(this.poiManager, p_203002_); // Neo: Must be called for all chunk unloading. Not just LevelChunks.
++                    this.chunkTypeCache.remove(p_203002_.getPos().toLong()); // Neo: Prevent chunk type cache from permanently retaining data for unloaded chunks
                      if (p_203002_ instanceof LevelChunk) {
                          ((LevelChunk)p_203002_).setLoaded(false);
 +                        net.neoforged.neoforge.common.NeoForge.EVENT_BUS.post(new net.neoforged.neoforge.event.level.ChunkEvent.Unload(p_203002_));
@@ -76,3 +80,24 @@
              EntityType<?> entitytype = p_140200_.getType();
              int i = entitytype.clientTrackingRange() * 16;
              if (i != 0) {
+@@ -1397,5 +_,20 @@
+                 this.updatePlayer(serverplayer);
+             }
+         }
++    }
++
++    /**
++     * Neo: PR #937
++     * This is for mainly pre-generation usage such as Neoforge's generate command.
++     * Use this to schedule chunk load tasks into ChunkTaskPriorityQueueSorter so a chunk is fully finished all of their tasks before scheduling more chunks to load.
++     * Reason for this is when scheduling a huge ton of Full Status chunk tasks to the server (to load chunks),
++     * you could cause the server to only process those loading tasks and never reach the two chunk tasks that are
++     * automatically scheduled to run after the chunk is loaded to Full. As a result of flooding the system with Full Status chunk tasks,
++     * the queue for the two kind of successor chunk tasks will grow and become a memory leak of lambdas and chunk references.
++     * Use this method to schedule tasks for loading chunks in your whenCompleteAsync method call so the tasks gets processed properly over time and not leak.
++     * See {@link net.neoforged.neoforge.server.command.generation.GenerationTask#enqueueChunks} as an example usage of this method.
++     */
++    public void scheduleOnMainThreadMailbox(ChunkTaskPriorityQueueSorter.Message<Runnable> msg) {
++        mainThreadMailbox.tell(msg);
+     }
+ }

patches/net/minecraft/world/level/chunk/storage/SectionStorage.java.patch

@@ -0,0 +1,15 @@
+--- a/net/minecraft/world/level/chunk/storage/SectionStorage.java
++++ b/net/minecraft/world/level/chunk/storage/SectionStorage.java
+@@ -231,4 +_,12 @@
+     public void close() throws IOException {
+         this.simpleRegionStorage.close();
+     }
++
++    /**
++     * Neo: Removes the data for the given chunk position.
++     * See PR #937
++     */
++    public void remove(long sectionPosAsLong) {
++        this.storage.remove(sectionPosAsLong);
++    }
+ }

src/main/java/net/neoforged/neoforge/common/CommonHooks.java

@@ -44,6 +44,7 @@
 import net.minecraft.core.HolderLookup;
 import net.minecraft.core.HolderSet;
 import net.minecraft.core.Registry;
+import net.minecraft.core.SectionPos;
 import net.minecraft.core.component.DataComponentMap;
 import net.minecraft.core.component.DataComponents;
 import net.minecraft.core.particles.ParticleTypes;
@@ -89,6 +90,7 @@
 import net.minecraft.world.entity.ai.attributes.AttributeModifier;
 import net.minecraft.world.entity.ai.attributes.AttributeSupplier;
 import net.minecraft.world.entity.ai.attributes.DefaultAttributes;
+import net.minecraft.world.entity.ai.village.poi.PoiManager;
 import net.minecraft.world.entity.item.ItemEntity;
 import net.minecraft.world.entity.monster.EnderMan;
 import net.minecraft.world.entity.player.Player;
@@ -113,6 +115,7 @@
 import net.minecraft.world.item.enchantment.EnchantmentHelper;
 import net.minecraft.world.item.enchantment.Enchantments;
 import net.minecraft.world.item.enchantment.ItemEnchantments;
+import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.GameType;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.biome.Biome;
@@ -124,6 +127,7 @@
 import net.minecraft.world.level.block.entity.BlockEntity;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.block.state.pattern.BlockInWorld;
+import net.minecraft.world.level.chunk.ChunkAccess;
 import net.minecraft.world.level.gameevent.GameEvent;
 import net.minecraft.world.level.material.Fluid;
 import net.minecraft.world.level.material.FluidState;
@@ -1338,4 +1342,22 @@ private static boolean overridesEqualsAndHashCode(Class<?> clazz) {
             throw new RuntimeException("Failed to check for component equals and hashCode", exception);
         }
     }
+
+    /**
+     * The goal here is to fix the POI memory leak that happens due to
+     * {@link net.minecraft.world.level.chunk.storage.SectionStorage#storage} field never
+     * actually removing POIs long after they become irrelevant. We do it here in chunk unload event
+     * so that chunk that are fully unloaded now gets the POI removed from the POI cached storage map.
+     */
+    public static void onChunkUnload(PoiManager poiManager, ChunkAccess chunkAccess) {
+        ChunkPos chunkPos = chunkAccess.getPos();
+        poiManager.flush(chunkPos); // Make sure all POI in chunk are saved to disk first.
+
+        // Remove the cached POIs for this chunk's location.
+        int SectionPosMinY = SectionPos.blockToSectionCoord(chunkAccess.getMinBuildHeight());
+        for (int currentSectionY = 0; currentSectionY < chunkAccess.getSectionsCount(); currentSectionY++) {
+            long sectionPosKey = SectionPos.asLong(chunkPos.x, SectionPosMinY + currentSectionY, chunkPos.z);
+            poiManager.remove(sectionPosKey);
+        }
+    }
 }

src/main/java/net/neoforged/neoforge/server/command/GenerateCommand.java

@@ -34,7 +34,7 @@ class GenerateCommand {
 
         builder.then(Commands.literal("start")
                 .then(Commands.argument("pos", BlockPosArgument.blockPos())
-                        .then(Commands.argument("chunkRadius", IntegerArgumentType.integer(1, 1250)) // 20000 block radius limit
+                        .then(Commands.argument("chunkRadius", IntegerArgumentType.integer(1, 2500)) // 40000 block radius limit
                                 .then(Commands.argument("progressBar", BoolArgumentType.bool())
                                         .executes(ctx -> executeGeneration(ctx.getSource(), BlockPosArgument.getSpawnablePos(ctx, "pos"), getInt(ctx, "chunkRadius"), getBool(ctx, "progressBar"))))
                                 .executes(ctx -> executeGeneration(ctx.getSource(), BlockPosArgument.getSpawnablePos(ctx, "pos"), getInt(ctx, "chunkRadius"), true)))));

src/main/java/net/neoforged/neoforge/server/command/generation/GenerationTask.java

@@ -20,6 +20,7 @@
 import net.minecraft.server.level.ChunkHolder;
 import net.minecraft.server.level.ChunkMap;
 import net.minecraft.server.level.ChunkResult;
+import net.minecraft.server.level.ChunkTaskPriorityQueueSorter;
 import net.minecraft.server.level.ServerChunkCache;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.TicketType;
@@ -41,6 +42,7 @@ public class GenerationTask {
 
     private final MinecraftServer server;
     private final ServerChunkCache chunkSource;
+    private final ServerLevel serverLevel;
 
     private final Iterator<ChunkPos> iterator;
     private final int x;
@@ -64,6 +66,7 @@ public class GenerationTask {
     public GenerationTask(ServerLevel serverLevel, int x, int z, int radius) {
         this.server = serverLevel.getServer();
         this.chunkSource = serverLevel.getChunkSource();
+        this.serverLevel = serverLevel;
 
         this.iterator = new CoarseOnionIterator(radius, COARSE_CELL_SIZE);
         this.x = x;
@@ -154,14 +157,14 @@ private void enqueueChunks(LongList chunks) {
                 continue;
             }
 
-            holder.getOrScheduleFuture(ChunkStatus.FULL, chunkMap).whenComplete((result, throwable) -> {
+            holder.getOrScheduleFuture(ChunkStatus.FULL, chunkMap).whenCompleteAsync((result, throwable) -> {
                 if (throwable == null) {
                     this.acceptChunkResult(chunkLongPos, result);
                 } else {
                     LOGGER.warn("Encountered unexpected error while generating chunk", throwable);
                     this.acceptChunkResult(chunkLongPos, ChunkHolder.UNLOADED_CHUNK);
                 }
-            });
+            }, runnable -> chunkMap.scheduleOnMainThreadMailbox(ChunkTaskPriorityQueueSorter.message(holder, runnable)));
         }
     }
 
@@ -180,6 +183,13 @@ private void acceptChunkResult(long chunk, ChunkResult<ChunkAccess> result) {
         if (queuedCount <= QUEUE_THRESHOLD) {
             this.tryEnqueueTasks();
         }
+
+        // Help make sure pregen progress does not get completely lost if game crashes/shuts down before pregen is finished.
+        if (((this.okCount.get() + this.errorCount.get()) % 1000) == 999) {
+            this.server.submit(() -> {
+                this.serverLevel.save(null, false, true);
+            });
+        }
     }
 
     private LongList collectChunks(int count) {
@@ -209,7 +219,7 @@ private void acquireChunk(long chunk) {
 
     private void releaseChunk(long chunk) {
         ChunkPos pos = new ChunkPos(chunk);
-        this.chunkSource.addRegionTicket(NEOFORGE_GENERATE_FORCED, pos, 0, pos);
+        this.chunkSource.removeRegionTicket(NEOFORGE_GENERATE_FORCED, pos, 0, pos);
     }
 
     private boolean isChunkFullyGenerated(ChunkPos chunkPosInLocalSpace) {