Paper/patches/server/1036-Write-SavedData-IO-async.patch

From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Cryptite <cryptite@gmail.com>
Date: Tue, 27 Jun 2023 11:35:52 -0500
Subject: [PATCH] Write SavedData IO async

Co-Authored-By: Shane Freeder <theboyetronic@gmail.com>

diff --git a/src/main/java/io/papermc/paper/world/ThreadedWorldUpgrader.java b/src/main/java/io/papermc/paper/world/ThreadedWorldUpgrader.java
index e049fbe4038aaea896f45b11ce9ce8f05922c898..7f6d1ccd147e5593412567bb2934ce5662da7ef0 100644
--- a/src/main/java/io/papermc/paper/world/ThreadedWorldUpgrader.java
+++ b/src/main/java/io/papermc/paper/world/ThreadedWorldUpgrader.java
@@ -110,6 +110,15 @@ public class ThreadedWorldUpgrader {
             }
 
             this.threadPool.execute(new ConvertTask(info, regionPos.x >> 5, regionPos.z >> 5));
+            // Paper start - Write SavedData IO async
+            this.threadPool.execute(() -> {
+                try {
+                    worldPersistentData.close();
+                } catch (IOException exception) {
+                    LOGGER.error("Failed to close persistent world data", exception);
+                }
+            });
+            // Paper end - Write SavedData IO async
         }
         this.threadPool.shutdown();
 
diff --git a/src/main/java/net/minecraft/server/level/ServerChunkCache.java b/src/main/java/net/minecraft/server/level/ServerChunkCache.java
index 36caf354634d6675a3f1ec6829f4778e1d0623bc..b99f50604bafecbc68835974c9ed0caa91911a40 100644
--- a/src/main/java/net/minecraft/server/level/ServerChunkCache.java
+++ b/src/main/java/net/minecraft/server/level/ServerChunkCache.java
@@ -447,6 +447,13 @@ public class ServerChunkCache extends ChunkSource {
 
     public void close(boolean save) { // Paper - rewrite chunk system
         this.level.chunkTaskScheduler.chunkHolderManager.close(save, true); // Paper - rewrite chunk system
+        // Paper start - Write SavedData IO async
+        try {
+            this.dataStorage.close();
+        } catch (IOException exception) {
+            LOGGER.error("Failed to close persistent world data", exception);
+        }
+        // Paper end - Write SavedData IO async
     }
 
     // CraftBukkit start - modelled on below
diff --git a/src/main/java/net/minecraft/server/level/ServerLevel.java b/src/main/java/net/minecraft/server/level/ServerLevel.java
index 0b0fd9e32d2772098d1aade83083b38cac01b76e..0981d440d0dbfe4df668d1f3f1b5706a93bc4434 100644
--- a/src/main/java/net/minecraft/server/level/ServerLevel.java
+++ b/src/main/java/net/minecraft/server/level/ServerLevel.java
@@ -1492,7 +1492,7 @@ public class ServerLevel extends Level implements WorldGenLevel {
 
         try (co.aikar.timings.Timing ignored = this.timings.worldSave.startTiming()) {
             if (doFull) {
-                this.saveLevelData();
+                this.saveLevelData(true); // Paper - Write SavedData IO async
             }
 
             this.timings.worldSaveChunks.startTiming(); // Paper
@@ -1528,7 +1528,7 @@ public class ServerLevel extends Level implements WorldGenLevel {
                 progressListener.progressStartNoAbort(Component.translatable("menu.savingLevel"));
             }
 
-            this.saveLevelData();
+            this.saveLevelData(!close); // Paper - Write SavedData IO async
             if (progressListener != null) {
                 progressListener.progressStage(Component.translatable("menu.savingChunks"));
             }
@@ -1551,12 +1551,12 @@ public class ServerLevel extends Level implements WorldGenLevel {
         // CraftBukkit end
     }
 
-    private void saveLevelData() {
+    private void saveLevelData(boolean async) { // Paper - Write SavedData IO async
         if (this.dragonFight != null) {
             this.serverLevelData.setEndDragonFightData(this.dragonFight.saveData()); // CraftBukkit
         }
 
-        this.getChunkSource().getDataStorage().save();
+        this.getChunkSource().getDataStorage().save(async); // Paper - Write SavedData IO async
     }
 
     public <T extends Entity> List<? extends T> getEntities(EntityTypeTest<Entity, T> filter, Predicate<? super T> predicate) {
diff --git a/src/main/java/net/minecraft/util/worldupdate/WorldUpgrader.java b/src/main/java/net/minecraft/util/worldupdate/WorldUpgrader.java
index 639f72618a7c22fa94effa9d0406b97fffc64cb5..3e582c49069f2a820ba3baac03917493877d9875 100644
--- a/src/main/java/net/minecraft/util/worldupdate/WorldUpgrader.java
+++ b/src/main/java/net/minecraft/util/worldupdate/WorldUpgrader.java
@@ -116,7 +116,13 @@ public class WorldUpgrader {
         (new WorldUpgrader.PoiUpgrader(this)).upgrade();
         WorldUpgrader.LOGGER.info("Upgrading blocks");
         (new WorldUpgrader.ChunkUpgrader()).upgrade();
-        this.overworldDataStorage.save();
+        // Paper start - Write SavedData IO async
+        try {
+            this.overworldDataStorage.close();
+        } catch (final IOException e) {
+            LOGGER.error("Failed to close persistent world data", e);
+        }
+        // Paper end - Write SavedData IO async
         i = Util.getMillis() - i;
         WorldUpgrader.LOGGER.info("World optimizaton finished after {} seconds", i / 1000L);
         this.finished = true;
diff --git a/src/main/java/net/minecraft/world/level/saveddata/SavedData.java b/src/main/java/net/minecraft/world/level/saveddata/SavedData.java
index 9cc3850bb70dfbcf342651360314a19fd9ea3ecc..4cbb943b702baaeb8bfd2b558cc848e719cf095d 100644
--- a/src/main/java/net/minecraft/world/level/saveddata/SavedData.java
+++ b/src/main/java/net/minecraft/world/level/saveddata/SavedData.java
@@ -30,20 +30,36 @@ public abstract class SavedData {
         return this.dirty;
     }
 
+    // Paper start - Write SavedData IO async - joining is evil, but we assume the old blocking behavior here just for safety
+    @io.papermc.paper.annotation.DoNotUse
     public void save(File file, HolderLookup.Provider registryLookup) {
+        save(file, registryLookup, null).join();
+    }
+
+    public java.util.concurrent.CompletableFuture<Void> save(File file, HolderLookup.Provider registryLookup, @org.jetbrains.annotations.Nullable java.util.concurrent.ExecutorService ioExecutor) {
+        // Paper end - Write SavedData IO async
         if (this.isDirty()) {
             CompoundTag compoundTag = new CompoundTag();
             compoundTag.put("data", this.save(new CompoundTag(), registryLookup));
             NbtUtils.addCurrentDataVersion(compoundTag);
 
+            Runnable writeRunnable = () -> { // Paper - Write SavedData IO async
             try {
                 NbtIo.writeCompressed(compoundTag, file.toPath());
             } catch (IOException var5) {
                 LOGGER.error("Could not save data {}", this, var5);
             }
+            }; // Paper - Write SavedData IO async
 
             this.setDirty(false);
+            // Paper start - Write SavedData IO async
+            if (ioExecutor == null) {
+                return java.util.concurrent.CompletableFuture.runAsync(writeRunnable); // No executor, just use common pool
+            }
+            return java.util.concurrent.CompletableFuture.runAsync(writeRunnable, ioExecutor);
         }
+        return java.util.concurrent.CompletableFuture.completedFuture(null);
+        // Paper end - Write SavedData IO async
     }
 
     public static record Factory<T extends SavedData>(
diff --git a/src/main/java/net/minecraft/world/level/storage/DimensionDataStorage.java b/src/main/java/net/minecraft/world/level/storage/DimensionDataStorage.java
index 7d98d2911e9d6ec429ce7df1f1f2650c7ea32325..6e23e69abd56eeda3b52a22019e1b74ae10682e7 100644
--- a/src/main/java/net/minecraft/world/level/storage/DimensionDataStorage.java
+++ b/src/main/java/net/minecraft/world/level/storage/DimensionDataStorage.java
@@ -23,17 +23,19 @@ import net.minecraft.util.datafix.DataFixTypes;
 import net.minecraft.world.level.saveddata.SavedData;
 import org.slf4j.Logger;
 
-public class DimensionDataStorage {
+public class DimensionDataStorage implements java.io.Closeable { // Paper - Write SavedData IO async
     private static final Logger LOGGER = LogUtils.getLogger();
     public final Map<String, SavedData> cache = Maps.newHashMap();
     private final DataFixer fixerUpper;
     private final HolderLookup.Provider registries;
     private final File dataFolder;
+    protected final java.util.concurrent.ExecutorService ioExecutor; // Paper - Write SavedData IO async
 
     public DimensionDataStorage(File directory, DataFixer dataFixer, HolderLookup.Provider registryLookup) {
         this.fixerUpper = dataFixer;
         this.dataFolder = directory;
         this.registries = registryLookup;
+        this.ioExecutor = java.util.concurrent.Executors.newSingleThreadExecutor(new com.google.common.util.concurrent.ThreadFactoryBuilder().setNameFormat("DimensionDataIO - " + dataFolder.getParent() + " - %d").setDaemon(true).build()); // Paper - Write SavedData IO async
     }
 
     private File getDataFile(String id) {
@@ -123,10 +125,23 @@ public class DimensionDataStorage {
         return bl;
     }
 
-    public void save() {
+    // Paper start - Write SavedData IO async
+    @Override
+    public void close() throws IOException {
+        save(false);
+        this.ioExecutor.shutdown();
+    }
+    // Paper end - Write SavedData IO async
+
+    public void save(boolean async) { // Paper - Write SavedData IO async
         this.cache.forEach((id, state) -> {
             if (state != null) {
-                state.save(this.getDataFile(id), this.registries);
+                // Paper start - Write SavedData IO async
+                final java.util.concurrent.CompletableFuture<Void> save = state.save(this.getDataFile(id), this.registries, this.ioExecutor);
+                if (!async) {
+                    save.join();
+                }
+                // Paper end - Write SavedData IO async
             }
         });
     }