Move SavedData.save to IO worker to avoid synchronous IO on main thre…

…ad (#1796)
neoforged · Feb 12, 2025 · 1acab57 · 1acab57
1 parent 2bf49e4
commit 1acab57
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 4 deletions.
diff --git a/patches/net/minecraft/server/level/ServerLevel.java.patch b/patches/net/minecraft/server/level/ServerLevel.java.patch
@@ -112,11 +112,14 @@
          this.getProfiler().pop();
 
          for (Entity entity : p_8648_.getPassengers()) {
-@@ -807,6 +_,7 @@
+@@ -807,6 +_,10 @@
              } else {
                  this.entityManager.autoSave();
              }
 +            net.neoforged.neoforge.common.NeoForge.EVENT_BUS.post(new net.neoforged.neoforge.event.level.LevelEvent.Save(this));
++            if (p_8645_) {
++                net.neoforged.neoforge.common.IOUtilities.waitUntilIOWorkerComplete();
++            }
          }
      }
 

diff --git a/patches/net/minecraft/util/worldupdate/WorldUpgrader.java.patch b/patches/net/minecraft/util/worldupdate/WorldUpgrader.java.patch
@@ -0,0 +1,10 @@
+--- a/net/minecraft/util/worldupdate/WorldUpgrader.java
++++ b/net/minecraft/util/worldupdate/WorldUpgrader.java
+@@ -111,6 +_,7 @@
+         LOGGER.info("Upgrading blocks");
+         new WorldUpgrader.ChunkUpgrader().upgrade();
+         this.overworldDataStorage.save();
++        net.neoforged.neoforge.common.IOUtilities.waitUntilIOWorkerComplete();
+         i = Util.getMillis() - i;
+         LOGGER.info("World optimizaton finished after {} seconds", i / 1000L);
+         this.finished = true;
diff --git a/patches/net/minecraft/world/level/saveddata/SavedData.java.patch b/patches/net/minecraft/world/level/saveddata/SavedData.java.patch
@@ -1,15 +1,23 @@
 --- a/net/minecraft/world/level/saveddata/SavedData.java
 +++ b/net/minecraft/world/level/saveddata/SavedData.java
-@@ -37,7 +_,7 @@
+@@ -36,18 +_,26 @@
+             compoundtag.put("data", this.save(new CompoundTag(), p_324088_));
              NbtUtils.addCurrentDataVersion(compoundtag);
 
++            //copy the contents to handle mods that just append an existing Tag to the Compound
++            CompoundTag copied = compoundtag.copy();
++
++            net.neoforged.neoforge.common.IOUtilities.withIOWorker(()->{
              try {
 -                NbtIo.writeCompressed(compoundtag, p_77758_.toPath());
-+                net.neoforged.neoforge.common.IOUtilities.writeNbtCompressed(compoundtag, p_77758_.toPath());
++                net.neoforged.neoforge.common.IOUtilities.writeNbtCompressed(copied, p_77758_.toPath());
              } catch (IOException ioexception) {
                  LOGGER.error("Could not save data {}", this, ioexception);
              }
-@@ -47,7 +_,10 @@
++            });
+
+             this.setDirty(false);
+         }
      }
 
      public static record Factory<T extends SavedData>(

diff --git a/src/main/java/net/neoforged/neoforge/common/IOUtilities.java b/src/main/java/net/neoforged/neoforge/common/IOUtilities.java
@@ -17,8 +17,10 @@
 import java.nio.file.StandardCopyOption;
 import java.nio.file.StandardOpenOption;
 import java.nio.file.attribute.BasicFileAttributes;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
+import net.minecraft.Util;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.NbtIo;
 import org.apache.commons.io.output.CloseShieldOutputStream;
@@ -34,6 +36,8 @@ public final class IOUtilities {
             StandardOpenOption.TRUNCATE_EXISTING
     };
 
+    private static CompletableFuture<Void> saveDataTasks = CompletableFuture.completedFuture(null);
+
     private IOUtilities() {}
 
     /**
@@ -152,6 +156,15 @@ public static void atomicWrite(Path targetPath, WriteCallback writeCallback) thr
         }
     }
 
+    public static void withIOWorker(Runnable task) {
+        saveDataTasks = saveDataTasks.thenRunAsync(task, Util.ioPool());
+    }
+
+    public static void waitUntilIOWorkerComplete() {
+        saveDataTasks.join();
+        saveDataTasks = CompletableFuture.completedFuture(null);
+    }
+
     /**
      * Declares an interface which is functionally equivalent to {@link Consumer},
      * except supports the ability to throw IOExceptions that may occur.

diff --git a/tests/src/junit/java/net/neoforged/neoforge/unittest/IOUtilitiesWorkerTest.java b/tests/src/junit/java/net/neoforged/neoforge/unittest/IOUtilitiesWorkerTest.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.unittest;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import net.neoforged.neoforge.common.IOUtilities;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class IOUtilitiesWorkerTest {
+    /**
+     * Tests that waiting on the IO worker will complete remaining tasks
+     */
+    @Test
+    void testWaitOnWorker() {
+        AtomicBoolean value = new AtomicBoolean(false);
+        IOUtilities.withIOWorker(() -> {
+            try {
+                Thread.sleep(2000);
+            } catch (InterruptedException ignored) {}
+            value.set(true);
+        });
+        Assertions.assertFalse(value.get(), "Value should not be set yet");
+        IOUtilities.waitUntilIOWorkerComplete();
+        Assertions.assertTrue(value.get(), "Value should be set after waiting");
+    }
+}