From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: violetc <58360096+s-yh-china@users.noreply.github.com>
Date: Thu, 25 Jan 2024 01:16:49 +0800
Subject: [PATCH] Fast resume


diff --git a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkHolderManager.java b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkHolderManager.java
index b133b977bc7d452b5032809f84f8ac2ff96ae5bb..be8280d421a483b354f88d60117002181abe570b 100644
--- a/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkHolderManager.java
+++ b/src/main/java/io/papermc/paper/chunk/system/scheduling/ChunkHolderManager.java
@@ -566,6 +566,56 @@ public final class ChunkHolderManager {
         }
     }
 
+    // Leaves start - add custom ticket
+    public <T> boolean addTicketAtLevelCustom(final Ticket<T> ticket, final long chunk, final boolean lock) {
+        final long removeDelay = ticket.removeDelay;
+        if (ticket.getTicketLevel() > MAX_TICKET_LEVEL) {
+            return false;
+        }
+
+        final int chunkX = CoordinateUtils.getChunkX(chunk);
+        final int chunkZ = CoordinateUtils.getChunkZ(chunk);
+        final RegionFileIOThread.ChunkCoordinate chunkCoord = new RegionFileIOThread.ChunkCoordinate(chunk);
+
+        final ReentrantAreaLock.Node ticketLock = lock ? this.ticketLockArea.lock(chunkX, chunkZ) : null;
+        try {
+            final SortedArraySet<Ticket<?>> ticketsAtChunk = this.tickets.computeIfAbsent(chunkCoord, (final RegionFileIOThread.ChunkCoordinate keyInMap) -> {
+                return SortedArraySet.create(4);
+            });
+
+            final int levelBefore = getTicketLevelAt(ticketsAtChunk);
+            final Ticket<T> current = (Ticket<T>)ticketsAtChunk.replace(ticket);
+            final int levelAfter = getTicketLevelAt(ticketsAtChunk);
+
+            if (current != ticket) {
+                final long oldRemoveDelay = current.removeDelay;
+                if (removeDelay != oldRemoveDelay) {
+                    if (oldRemoveDelay != NO_TIMEOUT_MARKER && removeDelay == NO_TIMEOUT_MARKER) {
+                        this.removeExpireCount(chunkX, chunkZ);
+                    } else if (oldRemoveDelay == NO_TIMEOUT_MARKER) {
+                        // since old != new, we have that NO_TIMEOUT_MARKER != new
+                        this.addExpireCount(chunkX, chunkZ);
+                    }
+                }
+            } else {
+                if (removeDelay != NO_TIMEOUT_MARKER) {
+                    this.addExpireCount(chunkX, chunkZ);
+                }
+            }
+
+            if (levelBefore != levelAfter) {
+                this.updateTicketLevel(chunk, levelAfter);
+            }
+
+            return current == ticket;
+        } finally {
+            if (ticketLock != null) {
+                this.ticketLockArea.unlock(ticketLock);
+            }
+        }
+    }
+    // Leaves end - add custom ticket
+
     public <T> boolean removeTicketAtLevel(final TicketType<T> type, final ChunkPos chunkPos, final int level, final T identifier) {
         return this.removeTicketAtLevel(type, CoordinateUtils.getChunkKey(chunkPos), level, identifier);
     }
diff --git a/src/main/java/net/minecraft/server/MinecraftServer.java b/src/main/java/net/minecraft/server/MinecraftServer.java
index 5d290f9335a510f4ddab7a78a96c0d8eedd4682e..502dfe1b2d030f9a2f484ced7d519342d58b79d5 100644
--- a/src/main/java/net/minecraft/server/MinecraftServer.java
+++ b/src/main/java/net/minecraft/server/MinecraftServer.java
@@ -669,6 +669,7 @@ public abstract class MinecraftServer extends ReentrantBlockableEventLoop<TickTa
 
         org.leavesmc.leaves.bot.ServerBot.loadAllBot(); // Leaves - load resident bot
 
+        org.leavesmc.leaves.util.TicketHelper.tryToLoadTickets(); // Leaves - load ticket
         this.server.enablePlugins(org.bukkit.plugin.PluginLoadOrder.POSTWORLD);
         if (io.papermc.paper.plugin.PluginInitializerManager.instance().pluginRemapper != null) io.papermc.paper.plugin.PluginInitializerManager.instance().pluginRemapper.pluginsEnabled(); // Paper - Remap plugins
         io.papermc.paper.command.brigadier.PaperCommands.INSTANCE.setValid(); // Paper - reset invalid state for event fire below
@@ -943,6 +944,7 @@ public abstract class MinecraftServer extends ReentrantBlockableEventLoop<TickTa
     // CraftBukkit end
 
     public void stopServer() {
+        org.leavesmc.leaves.util.TicketHelper.tryToSaveTickets(); // Leaves - save ticket
         // CraftBukkit start - prevent double stopping on multiple threads
         synchronized(this.stopLock) {
             if (this.hasStopped) return;
diff --git a/src/main/java/org/leavesmc/leaves/util/TicketHelper.java b/src/main/java/org/leavesmc/leaves/util/TicketHelper.java
new file mode 100644
index 0000000000000000000000000000000000000000..47154500a4375ee8bfd45d5dbe372425774e9d11
--- /dev/null
+++ b/src/main/java/org/leavesmc/leaves/util/TicketHelper.java
@@ -0,0 +1,179 @@
+package org.leavesmc.leaves.util;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import io.papermc.paper.chunk.system.RegionizedPlayerChunkLoader;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.registries.Registries;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.MinecraftServer;
+import net.minecraft.server.level.DistanceManager;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.Ticket;
+import net.minecraft.server.level.TicketType;
+import net.minecraft.util.SortedArraySet;
+import net.minecraft.world.level.ChunkPos;
+import net.minecraft.world.level.storage.LevelResource;
+import org.leavesmc.leaves.LeavesConfig;
+
+import java.io.BufferedReader;
+import java.io.BufferedWriter;
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.Set;
+
+public class TicketHelper {
+
+    private static final Set<TicketType<?>> NEED_SAVED = Set.of(TicketType.PLAYER, TicketType.PORTAL, TicketType.ENTITY_LOAD, TicketType.POI_LOAD);
+
+    public static void tryToLoadTickets() {
+        if (!LeavesConfig.fastResume) {
+            return;
+        }
+
+        File file = MinecraftServer.getServer().getWorldPath(LevelResource.ROOT).resolve("chunk_tickets.leaves.json").toFile();
+        if (file.isFile()) {
+            try (BufferedReader bfr = Files.newBufferedReader(file.toPath(), StandardCharsets.UTF_8)) {
+                JsonObject json = new Gson().fromJson(bfr, JsonObject.class);
+                loadSavedChunkTickets(json);
+                file.delete();
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        }
+    }
+
+    public static void tryToSaveTickets() {
+        if (!LeavesConfig.fastResume) {
+            return;
+        }
+
+        File file = MinecraftServer.getServer().getWorldPath(LevelResource.ROOT).resolve("chunk_tickets.leaves.json").toFile();
+        if (!file.isFile()) {
+            try {
+                file.createNewFile();
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        }
+        try (BufferedWriter bfw = Files.newBufferedWriter(file.toPath(), StandardCharsets.UTF_8)) {
+            bfw.write(new Gson().toJson(getSavedChunkTickets()));
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    public static void loadSavedChunkTickets(JsonObject json) {
+        MinecraftServer server = MinecraftServer.getServer();
+        for (String worldKey : json.keySet()) {
+            ServerLevel level = server.getLevel(ResourceKey.create(Registries.DIMENSION, new ResourceLocation(worldKey)));
+            if (level == null) {
+                continue;
+            }
+
+            DistanceManager chunkDistanceManager = level.getChunkSource().chunkMap.distanceManager;
+            for (JsonElement chunkElement : json.get(worldKey).getAsJsonArray()) {
+                JsonObject chunkJson = (JsonObject) chunkElement;
+                long chunkKey = chunkJson.get("key").getAsLong();
+
+                for (JsonElement ticketElement : chunkJson.get("tickets").getAsJsonArray()) {
+                    Ticket<?> ticket = tickFormJson((JsonObject) ticketElement);
+                    chunkDistanceManager.getChunkHolderManager().addTicketAtLevelCustom(ticket, chunkKey, true);
+                }
+            }
+        }
+    }
+
+    public static JsonObject getSavedChunkTickets() {
+        JsonObject json = new JsonObject();
+
+        for (ServerLevel level : MinecraftServer.getServer().getAllLevels()) {
+            JsonArray levelArray = new JsonArray();
+            DistanceManager chunkDistanceManager = level.getChunkSource().chunkMap.distanceManager;
+
+            for (Long2ObjectMap.Entry<SortedArraySet<Ticket<?>>> chunkTickets : chunkDistanceManager.getChunkHolderManager().getTicketsCopy().long2ObjectEntrySet()) {
+                long chunkKey = chunkTickets.getLongKey();
+                JsonArray ticketArray = new JsonArray();
+                SortedArraySet<Ticket<?>> tickets = chunkTickets.getValue();
+
+                for (Ticket<?> ticket : tickets) {
+                    if (!NEED_SAVED.contains(ticket.getType())) {
+                        continue;
+                    }
+
+                    ticketArray.add(ticketToJson(ticket));
+                }
+
+                if (!ticketArray.isEmpty()) {
+                    JsonObject chunkJson = new JsonObject();
+                    chunkJson.addProperty("key", chunkKey);
+                    chunkJson.add("tickets", ticketArray);
+                    levelArray.add(chunkJson);
+                }
+            }
+
+            if (!levelArray.isEmpty()) {
+                json.add(level.dimension().location().toString(), levelArray);
+            }
+        }
+
+        return json;
+    }
+
+    private static JsonObject ticketToJson(Ticket<?> ticket) {
+        JsonObject json = new JsonObject();
+        json.addProperty("type", ticket.getType().toString());
+        json.addProperty("ticketLevel", ticket.getTicketLevel());
+        json.addProperty("removeDelay", ticket.removeDelay);
+        if (ticket.key instanceof BlockPos pos) {
+            json.addProperty("key", pos.asLong());
+        } else if (ticket.key instanceof ChunkPos pos) {
+            json.addProperty("key", pos.toLong());
+        } else if (ticket.key instanceof Long l) {
+            json.addProperty("key", l);
+        }
+        return json;
+    }
+
+    private static <T> Ticket<T> tickFormJson(JsonObject json) {
+        TicketType<?> ticketType = null;
+        Object key = null;
+        switch (json.get("type").getAsString()) {
+            case "player" -> {
+                ticketType = TicketType.PLAYER;
+                key = new ChunkPos(json.get("key").getAsLong());
+            }
+            case "portal" -> {
+                ticketType = TicketType.PORTAL;
+                key = BlockPos.of(json.get("key").getAsLong());
+            }
+            case "entity_load" -> {
+                ticketType = TicketType.ENTITY_LOAD;
+                key = json.get("key").getAsLong();
+            }
+            case "poi_load" -> {
+                ticketType = TicketType.POI_LOAD;
+                key = json.get("key").getAsLong();
+            }
+            case "region_player_ticket" -> {
+                ticketType = RegionizedPlayerChunkLoader.REGION_PLAYER_TICKET;
+                key = json.get("key").getAsLong();
+            }
+        }
+
+        if (ticketType == null) {
+            throw new IllegalArgumentException("???");
+        }
+
+        int ticketLevel = json.get("ticketLevel").getAsInt();
+        long removeDelay = json.get("removeDelay").getAsLong();
+
+        return new Ticket<T>((TicketType<T>) ticketType, ticketLevel, (T) key, removeDelay);
+    }
+}