From 2c97e8dc6edbec887a8dd225c753db2188d62bce Mon Sep 17 00:00:00 2001
From: Spottedleaf <Spottedleaf@users.noreply.github.com>
Date: Fri, 28 Jun 2024 18:08:46 -0700
Subject: [PATCH] Fairly execute chunk tasks for all worlds

To ensure that a single world does not eat up allocated
processing time for tasks, execute tasks for all worlds
before returning from pollTaskInternal().
---
 .../chunk_system/MinecraftServerMixin.java    | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/main/java/ca/spottedleaf/moonrise/mixin/chunk_system/MinecraftServerMixin.java b/src/main/java/ca/spottedleaf/moonrise/mixin/chunk_system/MinecraftServerMixin.java
index da42c68..aae0537 100644
--- a/src/main/java/ca/spottedleaf/moonrise/mixin/chunk_system/MinecraftServerMixin.java
+++ b/src/main/java/ca/spottedleaf/moonrise/mixin/chunk_system/MinecraftServerMixin.java
@@ -8,10 +8,13 @@ import ca.spottedleaf.moonrise.patches.chunk_system.server.ChunkSystemMinecraftS
 import net.minecraft.commands.CommandSource;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.ServerInfo;
+import net.minecraft.server.ServerTickRateManager;
 import net.minecraft.server.TickTask;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.util.thread.ReentrantBlockableEventLoop;
+import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Overwrite;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.Unique;
 import org.spongepowered.asm.mixin.injection.At;
@@ -19,6 +22,7 @@ import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+import java.util.Iterator;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
@@ -32,6 +36,12 @@ public abstract class MinecraftServerMixin extends ReentrantBlockableEventLoop<T
     @Shadow
     public abstract boolean saveAllChunks(boolean bl, boolean bl2, boolean bl3);
 
+    @Shadow
+    @Final
+    private ServerTickRateManager tickRateManager;
+
+    @Shadow
+    protected abstract boolean haveTime();
 
     public MinecraftServerMixin(String string) {
         super(string);
@@ -45,6 +55,31 @@ public abstract class MinecraftServerMixin extends ReentrantBlockableEventLoop<T
         this.chunkSystemCrash = throwable;
     }
 
+    /**
+     * @reason Force execution of tasks for all worlds, so that the first world does not hog all of the task processing
+     * @author Spottedleaf
+     */
+    @Overwrite
+    private boolean pollTaskInternal() {
+        if (super.pollTask()) {
+            return true;
+        }
+
+        if (this.tickRateManager.isSprinting() || this.haveTime()) {
+            boolean ret = false;
+            for (final ServerLevel world : this.getAllLevels()) {
+                if (world.getChunkSource().pollTask()) {
+                    ret = true;
+                }
+            }
+
+            return ret;
+        }
+
+        return false;
+    }
+
+
     /**
      * @reason Force response to chunk system crash
      * @author Spottedleaf