Drastically simplify and document chunk system memory usage patch

Author: embeddedt

src/main/java/org/embeddedt/modernfix/common/mixin/bugfix/paper_chunk_patches/ChunkMapMixin.java

@@ -1,22 +1,18 @@
 package org.embeddedt.modernfix.common.mixin.bugfix.paper_chunk_patches;
 
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 import com.mojang.datafixers.util.Either;
 import net.minecraft.server.level.*;
-import net.minecraft.server.level.progress.ChunkProgressListener;
 import net.minecraft.util.thread.BlockableEventLoop;
-import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.chunk.ChunkAccess;
 import net.minecraft.world.level.chunk.ChunkStatus;
-import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplateManager;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
-import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.ModifyArg;
-import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
-import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 
@@ -25,18 +21,6 @@
 public abstract class ChunkMapMixin {
     @Shadow @Final private BlockableEventLoop<Runnable> mainThreadExecutor;
 
-    @Shadow @Final private ChunkMap.DistanceManager distanceManager;
-
-    @Shadow protected abstract CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>> protoChunkToFullChunk(ChunkHolder arg);
-
-    @Shadow @Final private ServerLevel level;
-    @Shadow @Final private ThreadedLevelLightEngine lightEngine;
-    @Shadow @Final private ChunkProgressListener progressListener;
-
-    @Shadow protected abstract CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>> scheduleChunkGeneration(ChunkHolder chunkHolder, ChunkStatus chunkStatus);
-
-    @Shadow @Final private StructureTemplateManager structureTemplateManager;
-
     /* https://github.com/PaperMC/Paper/blob/ver/1.17.1/patches/server/0752-Fix-chunks-refusing-to-unload-at-low-TPS.patch */
     @ModifyArg(method = "prepareAccessibleChunk", at = @At(value = "INVOKE", target = "Ljava/util/concurrent/CompletableFuture;thenApplyAsync(Ljava/util/function/Function;Ljava/util/concurrent/Executor;)Ljava/util/concurrent/CompletableFuture;"), index = 1)
     private Executor useMainThreadExecutor(Executor executor) {
@@ -45,31 +29,24 @@ private Executor useMainThreadExecutor(Executor executor) {
 
     /**
      * @author embeddedt
-     * @reason revert 1.17 chunk system changes, significantly reduces time and RAM needed to load chunks
+     * @reason 1.17+ uses getNow to check if the parent future is ready, and calls scheduleChunkGeneration as soon as
+     * it is found to not be ready. In the latter scenario, a massive number of extra CompletableFutures are allocated
+     * even if they are not actually necessary if the future is waited for. To prevent this, if the parent future
+     * is not done, we wait for it to complete and then retry schedule(). This will either detect an adequate
+     * status and return a loading future, or re-enter this injector with the parent future completed, in which case
+     * we proceed to schedule generation as originally requested.
      */
-    @Inject(method = "schedule", at = @At("HEAD"), cancellable = true)
-    private void useLegacySchedulingLogic(ChunkHolder holder, ChunkStatus requiredStatus, CallbackInfoReturnable<CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>>> cir) {
-        if(requiredStatus != ChunkStatus.EMPTY && !requiredStatus.hasLoadDependencies()) {
-            ChunkPos chunkpos = holder.getPos();
-            CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>> future = holder.getOrScheduleFuture(requiredStatus.getParent(), (ChunkMap)(Object)this);
-            cir.setReturnValue(future.thenComposeAsync((either) -> {
-                Optional<ChunkAccess> optional = either.left();
-
-                if (requiredStatus == ChunkStatus.LIGHT) {
-                    this.distanceManager.addTicket(TicketType.LIGHT, chunkpos, 33 + ChunkStatus.getDistance(ChunkStatus.LIGHT), chunkpos);
-                }
-
-                // from original method
-                if (optional.isPresent() && optional.get().getStatus().isOrAfter(requiredStatus)) {
-                    CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>> completablefuture = requiredStatus.load(this.level, this.structureTemplateManager, this.lightEngine, (arg2) -> {
-                        return this.protoChunkToFullChunk(holder);
-                    }, (ChunkAccess)optional.get());
-                    this.progressListener.onStatusChange(chunkpos, requiredStatus);
-                    return completablefuture;
-                } else {
-                    return this.scheduleChunkGeneration(holder, requiredStatus);
-                }
-            }, this.mainThreadExecutor).thenComposeAsync(CompletableFuture::completedFuture, this.mainThreadExecutor));
+    @WrapOperation(method = "schedule", at = @At(value = "INVOKE", target = "Lnet/minecraft/server/level/ChunkMap;scheduleChunkGeneration(Lnet/minecraft/server/level/ChunkHolder;Lnet/minecraft/world/level/chunk/ChunkStatus;)Ljava/util/concurrent/CompletableFuture;"))
+    private CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>> mfix$avoidSchedulingGenerationPrematurely(ChunkMap map, ChunkHolder holder, ChunkStatus status, Operation<CompletableFuture<Either<ChunkAccess, ChunkHolder.ChunkLoadingFailure>>> original) {
+        if (!status.hasLoadDependencies()) {
+            var parentFuture = holder.getOrScheduleFuture(status.getParent(), map);
+            if (!parentFuture.isDone()) {
+                return parentFuture.thenComposeAsync(
+                        either -> map.schedule(holder, status),
+                        this.mainThreadExecutor
+                );
+            }
         }
+        return original.call(map, holder, status);
     }
 }