LuminolMC/patches/server/0030-Kaiiju-Entity-tick-and-removal-limiter.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: MrHua269 <wangxyper@163.com>
Date: Wed, 4 Dec 2024 23:39:18 +0800
Subject: [PATCH] Kaiiju Entity tick and removal limiter


diff --git a/build.gradle.kts b/build.gradle.kts
index 8f23bf19618382ccf5fd10a0b17b57cd445dea58..4b79f96103c98896332113ffe5f0e22cf08ffdd1 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -27,6 +27,7 @@ abstract class MockitoAgentProvider : CommandLineArgumentProvider {
 dependencies {
     implementation(project(":luminol-api")) // Folia // Luminol
     implementation("com.electronwill.night-config:toml:3.6.6") // Luminol - Night config
+    implementation("io.github.classgraph:classgraph:4.8.158") // Kaiiju - Entity throttling & Removal
     implementation("ca.spottedleaf:concurrentutil:0.0.2") // Paper - Add ConcurrentUtil dependency
     // Paper start
     implementation("org.jline:jline-terminal-ffm:3.27.1") // use ffm on java 22+
diff --git a/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityLimits.java b/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityLimits.java
new file mode 100644
index 0000000000000000000000000000000000000000..40e80fc685a42bbaeea3e6e64754121178cc7d22
--- /dev/null
+++ b/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityLimits.java
@@ -0,0 +1,141 @@
+package dev.kaiijumc.kaiiju;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.logging.Level;
+
+import com.google.common.base.Throwables;
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
+import io.github.classgraph.ClassGraph;
+import io.github.classgraph.ClassInfo;
+import io.github.classgraph.ScanResult;
+import org.slf4j.Logger;
+
+import com.mojang.logging.LogUtils;
+import net.minecraft.world.entity.Entity;
+import org.bukkit.Bukkit;
+import org.bukkit.configuration.InvalidConfigurationException;
+import org.bukkit.configuration.file.YamlConfiguration;
+
+@SuppressWarnings("unused")
+public class KaiijuEntityLimits {
+    private static final Logger LOGGER = LogUtils.getLogger();
+    private static final File CONFIG_FOLDER = new File("luminol_config");
+
+    protected static final String HEADER =
+        "Per region entity limits for Kaiiju.\n"
+        + "If there are more of particular entity type in a region than limit, entity ticking will be throttled.\n"
+        + "Example: for Wither limit 100 & 300 Withers in a region -> 100 Withers tick every tick & every Wither ticks every 3 ticks.\n"
+        + "Available entities: GlowSquid, Ambient, Bat, Animal, Bee, Cat, Chicken, Cod, Cow, Dolphin, Fish, FishSchool, Fox, Golem, IronGolem, "
+        + "MushroomCow, Ocelot, Panda, Parrot, Perchable, Pig, PolarBear, PufferFish, Rabbit, Salmon, Sheep, Snowman, Squid, TropicalFish, Turtle, "
+        + "WaterAnimal, Wolf, Allay, Axolotl, Camel, Frog, Tadpole, Goat, Horse, HorseAbstract, HorseChestedAbstract, HorseDonkey, HorseMule, "
+        + "HorseSkeleton, HorseZombie, Llama, LlamaTrader, Sniffer, EnderCrystal, EnderDragon, Wither, ArmorStand, Hanging, ItemFrame, Leash, "
+        + "Painting, GlowItemFrame, FallingBlock, Item, TNTPrimed, Blaze, CaveSpider, Creeper, Drowned, Enderman, Endermite, Evoker, Ghast, "
+        + "GiantZombie, Guardian, GuardianElder, IllagerAbstract, IllagerIllusioner, IllagerWizard, MagmaCube, Monster, MonsterPatrolling, Phantom, "
+        + "PigZombie, Pillager, Ravager, Shulker, Silverfish, Skeleton, SkeletonAbstract, SkeletonStray, SkeletonWither, Slime, Spider, Strider, Vex, "
+        + "Vindicator, Witch, Zoglin, Zombie, ZombieHusk, ZombieVillager, Hoglin, Piglin, PiglinAbstract, PiglinBrute, Warden, Villager, "
+        + "VillagerTrader, Arrow, DragonFireball, Egg, EnderPearl, EnderSignal, EvokerFangs, Fireball, FireballFireball, Fireworks, FishingHook, "
+        + "LargeFireball, LlamaSpit, Potion, Projectile, ProjectileThrowable, ShulkerBullet, SmallFireball, Snowball, SpectralArrow, ThrownExpBottle, "
+        + "ThrownTrident, TippedArrow, WitherSkull, Raider, ChestBoat, Boat, MinecartAbstract, MinecartChest, MinecartCommandBlock, MinecartContainer, "
+        + "MinecartFurnace, MinecartHopper, MinecartMobSpawner, MinecartRideable, MinecartTNT\n";
+    protected static final File ENTITY_LIMITS_FILE = new File(CONFIG_FOLDER, "kaiiju_entity_limits.yml");
+    public static YamlConfiguration entityLimitsConfig;
+    public static boolean enabled = false;
+
+    protected static Map<Class<? extends Entity>, EntityLimit> entityLimits;
+
+    static final String ENTITY_PREFIX = "Entity";
+
+    public static void init() {
+        init(true);
+    }
+
+    private static void init(boolean setup) {
+        entityLimitsConfig = new YamlConfiguration();
+
+        if (ENTITY_LIMITS_FILE.exists()) {
+            try {
+                entityLimitsConfig.load(ENTITY_LIMITS_FILE);
+            } catch (InvalidConfigurationException ex) {
+                Bukkit.getLogger().log(Level.SEVERE, "Could not load kaiiju_entity_limits.yml, please correct your syntax errors", ex);
+                throw Throwables.propagate(ex);
+            } catch (IOException ignore) {}
+        } else {
+            if (setup) {
+                entityLimitsConfig.options().header(HEADER);
+                entityLimitsConfig.options().copyDefaults(true);
+                entityLimitsConfig.set("enabled", enabled);
+                entityLimitsConfig.set("Axolotl.limit", 1000);
+                entityLimitsConfig.set("Axolotl.removal", 2000);
+                try {
+                    entityLimitsConfig.save(ENTITY_LIMITS_FILE);
+                } catch (IOException ex) {
+                    Bukkit.getLogger().log(Level.SEVERE, "Could not save " + ENTITY_LIMITS_FILE, ex);
+                }
+            }
+        }
+
+        enabled = entityLimitsConfig.getBoolean("enabled");
+
+        entityLimits = new Object2ObjectOpenHashMap<>();
+        try (ScanResult scanResult = new ClassGraph().enableAllInfo().acceptPackages("net.minecraft.world.entity").scan()) {
+            Map<String, ClassInfo> entityClasses = new HashMap<>();
+            for (ClassInfo classInfo : scanResult.getAllClasses()) {
+                Class<?> entityClass = Class.forName(classInfo.getName());
+                if (Entity.class.isAssignableFrom(entityClass)) {
+                    String entityName = extractEntityName(entityClass.getSimpleName());
+                    entityClasses.put(entityName, classInfo);
+                }
+            }
+
+            for (String key : entityLimitsConfig.getKeys(false)) {
+                if (key.equals("enabled")) {
+                    continue;
+                }
+
+                if (!entityClasses.containsKey(key)) {
+                    LOGGER.error("Unknown entity '" + key + "' in kaiiju-entity-limits.yml, skipping");
+                    continue;
+                }
+                int limit = entityLimitsConfig.getInt(key + ".limit");
+                int removal = entityLimitsConfig.getInt(key + ".removal");
+
+                if (limit < 1) {
+                    LOGGER.error(key + " has a limit less than the minimum of 1, ignoring");
+                    continue;
+                }
+                if (removal <= limit && removal != -1) {
+                    LOGGER.error(key + " has a removal limit that is less than or equal to its limit, setting removal to limit * 10");
+                    removal = limit * 10;
+                }
+
+                entityLimits.put((Class<? extends Entity>) Class.forName(entityClasses.get(key).getName()), new EntityLimit(limit, removal));
+            }
+        } catch (ClassNotFoundException e) {
+            e.printStackTrace();
+        }
+    }
+
+    public static EntityLimit getEntityLimit(Entity entity) {
+        return entityLimits.get(entity.getClass());
+    }
+
+    private static String extractEntityName(String input) {
+        int prefixLength = ENTITY_PREFIX.length();
+
+        if (input.length() <= prefixLength || !input.startsWith(ENTITY_PREFIX)) {
+            return input;
+        } else {
+            return input.substring(prefixLength);
+        }
+    }
+
+    public record EntityLimit(int limit, int removal) {
+        @Override
+        public String toString() {
+            return "EntityLimit{limit=" + limit + ", removal=" + removal + "}";
+        }
+    }
+}
diff --git a/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityThrottler.java b/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityThrottler.java
new file mode 100644
index 0000000000000000000000000000000000000000..eb690efacf083e4ff3e321578b12c534e6a40196
--- /dev/null
+++ b/src/main/java/dev/kaiijumc/kaiiju/KaiijuEntityThrottler.java
@@ -0,0 +1,84 @@
+package dev.kaiijumc.kaiiju;
+
+import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
+import net.minecraft.world.entity.Entity;
+import io.papermc.paper.threadedregions.RegionizedWorldData;
+
+public class KaiijuEntityThrottler {
+    private static class TickInfo {
+        int currentTick;
+        int continueFrom;
+        int toTick;
+        int toRemove;
+    }
+
+    public static class EntityThrottlerReturn {
+        public boolean skip;
+        public boolean remove;
+    }
+
+    private final Object2ObjectOpenHashMap<KaiijuEntityLimits.EntityLimit, TickInfo> entityLimitTickInfoMap = new Object2ObjectOpenHashMap<>();
+
+    public void tickLimiterStart() {
+        for (TickInfo tickInfo : entityLimitTickInfoMap.values()) {
+            tickInfo.currentTick = 0;
+        }
+    }
+
+    public EntityThrottlerReturn tickLimiterShouldSkip(Entity entity) {
+        EntityThrottlerReturn retVal = new EntityThrottlerReturn();
+        if (entity.isRemoved()) return retVal;
+        KaiijuEntityLimits.EntityLimit entityLimit = KaiijuEntityLimits.getEntityLimit(entity);
+
+        if (entityLimit != null) {
+            TickInfo tickInfo = entityLimitTickInfoMap.computeIfAbsent(entityLimit, el -> {
+                TickInfo newTickInfo = new TickInfo();
+                newTickInfo.toTick = entityLimit.limit();
+                return newTickInfo;
+            });
+
+            tickInfo.currentTick++;
+            if (tickInfo.currentTick <= tickInfo.toRemove && entityLimit.removal() > 0) {
+                retVal.skip = false;
+                retVal.remove = true;
+                return retVal;
+            }
+
+            if (tickInfo.currentTick < tickInfo.continueFrom) {
+                retVal.skip = true;
+                return retVal;
+            }
+            if (tickInfo.currentTick - tickInfo.continueFrom < tickInfo.toTick) {
+                retVal.skip = false;
+                return retVal;
+            }
+            retVal.skip = true;
+            return retVal;
+        } else {
+            retVal.skip = false;
+            return retVal;
+        }
+    }
+
+    public void tickLimiterFinish(RegionizedWorldData regionizedWorldData) {
+        for (var entry : entityLimitTickInfoMap.entrySet()) {
+            KaiijuEntityLimits.EntityLimit entityLimit = entry.getKey();
+            TickInfo tickInfo = entry.getValue();
+
+            int additionals = 0;
+            int nextContinueFrom = tickInfo.continueFrom + tickInfo.toTick;
+            if (nextContinueFrom >= tickInfo.currentTick) {
+                additionals = entityLimit.limit() - (tickInfo.currentTick - tickInfo.continueFrom);
+                nextContinueFrom = 0;
+            }
+            tickInfo.continueFrom = nextContinueFrom;
+            tickInfo.toTick = entityLimit.limit() + additionals;
+
+            if (tickInfo.toRemove == 0 && tickInfo.currentTick > entityLimit.removal()) {
+                tickInfo.toRemove = tickInfo.currentTick - entityLimit.removal();
+            } else if (tickInfo.toRemove != 0) {
+                tickInfo.toRemove = 0;
+            }
+        }
+    }
+}
diff --git a/src/main/java/io/papermc/paper/threadedregions/RegionizedWorldData.java b/src/main/java/io/papermc/paper/threadedregions/RegionizedWorldData.java
index 1b741d4bccfd45beeec43300f44770516c0d850e..3c37ad27488486f9bb0f972369ccaee2284df673 100644
--- a/src/main/java/io/papermc/paper/threadedregions/RegionizedWorldData.java
+++ b/src/main/java/io/papermc/paper/threadedregions/RegionizedWorldData.java
@@ -354,6 +354,7 @@ public final class RegionizedWorldData {
     private final IteratorSafeOrderedReferenceSet<Mob> navigatingMobs = new IteratorSafeOrderedReferenceSet<>();
     public final ReferenceList<Entity> trackerEntities = new ReferenceList<>(EMPTY_ENTITY_ARRAY); // Moonrise - entity tracker
     public final ReferenceList<Entity> trackerUnloadedEntities = new ReferenceList<>(EMPTY_ENTITY_ARRAY); // Moonrise - entity tracker
+    public final dev.kaiijumc.kaiiju.KaiijuEntityThrottler entityThrottler = new dev.kaiijumc.kaiiju.KaiijuEntityThrottler(); // Kaiiju
 
     // block ticking
     private final ObjectLinkedOpenHashSet<BlockEventData> blockEvents = new ObjectLinkedOpenHashSet<>();
diff --git a/src/main/java/me/earthme/luminol/config/modules/misc/KaiijuEntityLimiterConfig.java b/src/main/java/me/earthme/luminol/config/modules/misc/KaiijuEntityLimiterConfig.java
new file mode 100644
index 0000000000000000000000000000000000000000..038d9ab60cfac7e40e7c0c0644fa0a0d035eac01
--- /dev/null
+++ b/src/main/java/me/earthme/luminol/config/modules/misc/KaiijuEntityLimiterConfig.java
@@ -0,0 +1,23 @@
+package me.earthme.luminol.config.modules.misc;
+
+import com.electronwill.nightconfig.core.file.CommentedFileConfig;
+import dev.kaiijumc.kaiiju.KaiijuEntityLimits;
+import me.earthme.luminol.config.EnumConfigCategory;
+import me.earthme.luminol.config.IConfigModule;
+
+public class KaiijuEntityLimiterConfig implements IConfigModule {
+    @Override
+    public EnumConfigCategory getCategory() {
+        return EnumConfigCategory.MISC;
+    }
+
+    @Override
+    public String getBaseName() {
+        return "kaiiju_entity_limiter";
+    }
+
+    @Override
+    public void onLoaded(CommentedFileConfig configInstance) {
+        KaiijuEntityLimits.init();
+    }
+}
diff --git a/src/main/java/net/minecraft/server/level/ServerLevel.java b/src/main/java/net/minecraft/server/level/ServerLevel.java
index 07037eb601f9dcae2cad5f3e3d5f9a0ac142b68a..5503d506c595296ecad09a3ce4497a365f216af5 100644
--- a/src/main/java/net/minecraft/server/level/ServerLevel.java
+++ b/src/main/java/net/minecraft/server/level/ServerLevel.java
@@ -811,6 +811,7 @@ public class ServerLevel extends Level implements ServerEntityGetter, WorldGenLe
             }
 
             profiler.startTimer(ca.spottedleaf.leafprofiler.LProfilerRegistry.ACTIVATE_ENTITIES); try { // Folia - profiler
+            if (dev.kaiijumc.kaiiju.KaiijuEntityLimits.enabled) regionizedWorldData.entityThrottler.tickLimiterStart(); // Kaiiju
             org.spigotmc.ActivationRange.activateEntities(this); // Spigot
             } finally { profiler.stopTimer(ca.spottedleaf.leafprofiler.LProfilerRegistry.ACTIVATE_ENTITIES); } // Folia - profiler
             profiler.startTimer(ca.spottedleaf.leafprofiler.LProfilerRegistry.ENTITY_TICK); try { // Folia - profiler
@@ -832,6 +833,13 @@ public class ServerLevel extends Level implements ServerEntityGetter, WorldGenLe
                                 entity.stopRiding();
                             }
 
+                            // Kaiiju start
+                            if (dev.kaiijumc.kaiiju.KaiijuEntityLimits.enabled) {
+                                dev.kaiijumc.kaiiju.KaiijuEntityThrottler.EntityThrottlerReturn throttle = regionizedWorldData.entityThrottler.tickLimiterShouldSkip(entity);
+                                if (throttle.remove && !entity.hasCustomName()) entity.remove(Entity.RemovalReason.DISCARDED);
+                                if (throttle.skip) return;
+                            }
+                            // Kaiiju end
                             gameprofilerfiller.push("tick");
                             this.guardEntityTick(this::tickNonPassenger, entity);
                             gameprofilerfiller.pop();
@@ -839,6 +847,7 @@ public class ServerLevel extends Level implements ServerEntityGetter, WorldGenLe
                     }
                 }
             });
+            if (dev.kaiijumc.kaiiju.KaiijuEntityLimits.enabled) regionizedWorldData.entityThrottler.tickLimiterFinish(regionizedWorldData); // Kaiiju
             } finally { profiler.stopTimer(ca.spottedleaf.leafprofiler.LProfilerRegistry.ENTITY_TICK); } // Folia - profiler
             gameprofilerfiller.pop();
             profiler.startTimer(ca.spottedleaf.leafprofiler.LProfilerRegistry.TILE_ENTITY); try { // Folia - profiler