Add system to mark chunks as "lazy unsaved".

Avoids serialization of chunks with computers in them every single tick.
Since computers state changes every tick they're running, flagging the chunk unsaved every chunk would be necessary. However, this is undesirable, because it triggers an NBT serialization ever tick. Which is slow.
To avoid this, we now track a "lazy unsaved" flag for chunks, which computers can set. This flag is applied to the real unsaved flag during explicit saves, right before the unsaved flag is checked. Explicit saves include chunk unloads, pausing in single-player, server stopping and the save command, for example.

build.gradle

@@ -5,6 +5,7 @@ buildscript {
     }
     dependencies {
         classpath group: 'net.minecraftforge.gradle', name: 'ForgeGradle', version: '5.1.+', changing: true
+        classpath group: 'org.spongepowered', name: 'mixingradle', version: '0.7.+'
     }
 }
 
@@ -15,6 +16,7 @@ plugins {
 }
 
 apply plugin: 'net.minecraftforge.gradle'
+apply plugin: 'org.spongepowered.mixin'
 apply from: 'minecraft.gradle'
 
 def getGitRef() {
@@ -60,6 +62,7 @@ repositories {
 
 dependencies {
     minecraft "net.minecraftforge:forge:${minecraft_version}-${forge_version}"
+    annotationProcessor 'org.spongepowered:mixin:0.8.4:processor'
 
     implementation "li.cil.sedna:sedna-${minecraft_version}-forge:1.0.10"
 
@@ -85,9 +88,7 @@ minecraft {
     mappings channel: 'official', version: minecraft_version
 
     runs {
-        client {
-            workingDirectory project.file('run')
-
+        all {
             property 'forge.logging.markers', 'REGISTRIES'
             property 'forge.logging.console.level', 'info'
 
@@ -96,38 +97,30 @@ minecraft {
                     source sourceSets.main
                 }
             }
+
+            arg "-mixin.config=oc2.mixins.json"
+        }
+
+        client {
+            workingDirectory project.file('run')
         }
 
         server {
             workingDirectory project.file('run')
-
-            property 'forge.logging.markers', 'REGISTRIES'
-            property 'forge.logging.console.level', 'info'
-
-            mods {
-                oc2 {
-                    source sourceSets.main
-                }
-            }
+            arg "--nogui"
         }
 
         data {
             workingDirectory project.file('run')
-
-            property 'forge.logging.markers', 'REGISTRIES'
-            property 'forge.logging.console.level', 'info'
-
             args '--mod', 'oc2', '--all', '--output', file('src/generated/resources/'), '--existing', file('src/main/resources')
-
-            mods {
-                oc2 {
-                    source sourceSets.main
-                }
-            }
         }
     }
 }
 
+mixin {
+    add sourceSets.main, "oc2.refmap.json"
+}
+
 task copyGeneratedResources(type: Copy) {
     from 'src/generated'
     into 'src/main'
@@ -145,7 +138,8 @@ jar {
             'Implementation-Title'    : project.name,
             'Implementation-Version'  : "${semver}",
             'Implementation-Vendor'   : 'Sangar',
-            'Implementation-Timestamp': new Date().format("yyyy-MM-dd'T'HH:mm:ssZ")
+            'Implementation-Timestamp': new Date().format("yyyy-MM-dd'T'HH:mm:ssZ"),
+            'MixinConfigs'            : 'oc2.mixins.json',
         ])
     }
 }

src/main/java/li/cil/oc2/common/blockentity/ComputerBlockEntity.java

@@ -66,6 +66,7 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
 
     private boolean hasAddedOwnDevices;
     private boolean isNeighborUpdateScheduled;
+    private LevelChunk chunk;
 
     ///////////////////////////////////////////////////////////////////
 
@@ -192,6 +193,10 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
             level.updateNeighborsAt(getBlockPos(), getBlockState().getBlock());
         }
 
+        // Just grab it again every tick, to avoid this becoming invalid if something tries to
+        // mess with this BlockEntity in unexpected ways.
+        chunk = level.getChunkAt(getBlockPos());
+
         virtualMachine.tick();
     }
 
@@ -299,14 +304,17 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
 
     ///////////////////////////////////////////////////////////////////
 
+    private <T> void sendToClientsTrackingComputer(final T message) {
+        if (chunk != null) {
+            Network.sendToClientsTrackingChunk(message, chunk);
+        }
+    }
+
+    ///////////////////////////////////////////////////////////////////
+
     private final class ComputerItemStackHandlers extends AbstractVMItemStackHandlers {
         public ComputerItemStackHandlers() {
-            super(
-                new GroupDefinition(DeviceTypes.MEMORY, MEMORY_SLOTS),
-                new GroupDefinition(DeviceTypes.HARD_DRIVE, HARD_DRIVE_SLOTS),
-                new GroupDefinition(DeviceTypes.FLASH_MEMORY, FLASH_MEMORY_SLOTS),
-                new GroupDefinition(DeviceTypes.CARD, CARD_SLOTS)
-            );
+            super(new GroupDefinition(DeviceTypes.MEMORY, MEMORY_SLOTS), new GroupDefinition(DeviceTypes.HARD_DRIVE, HARD_DRIVE_SLOTS), new GroupDefinition(DeviceTypes.FLASH_MEMORY, FLASH_MEMORY_SLOTS), new GroupDefinition(DeviceTypes.CARD, CARD_SLOTS));
         }
 
         @Override
@@ -387,13 +395,11 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
 
         @Override
         protected void sendTerminalUpdateToClient(final ByteBuffer output) {
-            Network.sendToClientsTrackingChunk(new ComputerTerminalOutputMessage(ComputerBlockEntity.this, output), virtualMachine.chunk);
+            sendToClientsTrackingComputer(new ComputerTerminalOutputMessage(ComputerBlockEntity.this, output));
         }
     }
 
     private final class ComputerVirtualMachine extends AbstractVirtualMachine {
-        private LevelChunk chunk;
-
         private ComputerVirtualMachine(final CommonDeviceBusController busController, final BaseAddressProvider baseAddressProvider) {
             super(busController);
             state.vmAdapter.setBaseAddressProvider(baseAddressProvider);
@@ -416,12 +422,8 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
         public void tick() {
             assert level != null;
 
-            if (chunk == null) {
-                chunk = level.getChunkAt(getBlockPos());
-            }
-
             if (isRunning()) {
-                chunk.setUnsaved(true);
+                ChunkUtils.setLazyUnsaved(level, getBlockPos());
             }
 
             super.tick();
@@ -445,8 +447,7 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
         protected void stopRunnerAndReset() {
             super.stopRunnerAndReset();
 
-            TerminalUtils.resetTerminal(terminal, output -> Network.sendToClientsTrackingChunk(
-                new ComputerTerminalOutputMessage(ComputerBlockEntity.this, output), chunk));
+            TerminalUtils.resetTerminal(terminal, output -> sendToClientsTrackingComputer(new ComputerTerminalOutputMessage(ComputerBlockEntity.this, output)));
         }
 
         @Override
@@ -456,7 +457,7 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
 
         @Override
         protected void handleBusStateChanged(final CommonDeviceBusController.BusState value) {
-            Network.sendToClientsTrackingChunk(new ComputerBusStateMessage(ComputerBlockEntity.this, value), chunk);
+            sendToClientsTrackingComputer(new ComputerBusStateMessage(ComputerBlockEntity.this, value));
 
             if (value == CommonDeviceBusController.BusState.READY && level != null) {
                 // Bus just became ready, meaning new devices may be available, meaning new
@@ -467,16 +468,12 @@ public final class ComputerBlockEntity extends ModBlockEntity implements Termina
 
         @Override
         protected void handleRunStateChanged(final VMRunState value) {
-            // This method can be called from disposal logic, so if we are disposed quickly enough
-            // chunk may not be initialized yet. Avoid resulting NRE in network logic.
-            if (chunk != null) {
-                Network.sendToClientsTrackingChunk(new ComputerRunStateMessage(ComputerBlockEntity.this, value), chunk);
-            }
+            sendToClientsTrackingComputer(new ComputerRunStateMessage(ComputerBlockEntity.this, value));
         }
 
         @Override
         protected void handleBootErrorChanged(@Nullable final Component value) {
-            Network.sendToClientsTrackingChunk(new ComputerBootErrorMessage(ComputerBlockEntity.this, value), chunk);
+            sendToClientsTrackingComputer(new ComputerBootErrorMessage(ComputerBlockEntity.this, value));
         }
     }
 }

src/main/java/li/cil/oc2/common/mixin/ChunkAccessMixin.java

@@ -0,0 +1,35 @@
+package li.cil.oc2.common.mixin;
+
+import li.cil.oc2.common.util.ChunkAccessExt;
+import li.cil.oc2.common.util.ChunkUtils;
+import net.minecraft.world.level.chunk.ChunkAccess;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+
+/**
+ * Tracks a "lazy unsaved" flag per {@link ChunkAccess} instance, to allow
+ * marking chunks as needing saving just in time for "hard" saves.
+ *
+ * @see ChunkMapMixin <c>ChunkMapMixin</c> for more information
+ * @see ChunkUtils
+ */
+@Mixin(ChunkAccess.class)
+public abstract class ChunkAccessMixin implements ChunkAccessExt {
+    @Shadow
+    protected volatile boolean unsaved;
+
+    private volatile boolean lazyUnsaved;
+
+    @Override
+    public void setLazyUnsaved() {
+        lazyUnsaved = true;
+    }
+
+    @Override
+    public void applyAndClearLazyUnsaved() {
+        if (!unsaved && lazyUnsaved) {
+            unsaved = true;
+        }
+        lazyUnsaved = false;
+    }
+}

src/main/java/li/cil/oc2/common/mixin/ChunkMapMixin.java

@@ -0,0 +1,96 @@
+package li.cil.oc2.common.mixin;
+
+import com.mojang.datafixers.DataFixer;
+import it.unimi.dsi.fastutil.longs.Long2ObjectLinkedOpenHashMap;
+import li.cil.oc2.common.util.ChunkAccessExt;
+import li.cil.oc2.common.util.ChunkUtils;
+import net.minecraft.server.level.ChunkHolder;
+import net.minecraft.server.level.ChunkMap;
+import net.minecraft.world.level.block.entity.BlockEntity;
+import net.minecraft.world.level.chunk.ChunkAccess;
+import net.minecraft.world.level.chunk.storage.ChunkStorage;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.nio.file.Path;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Hooks into {@link ChunkMap} saving code-paths for "hard" save operations.
+ * <p>
+ * Minecraft immediately serializes all chunk data, including {@link BlockEntity} NBT.
+ * This is a massive performance issue for blocks with state that changes every tick,
+ * such as computers and things accepting energy.
+ * <p>
+ * To avoid this per-frame serialization operations, we track a "lazy unsaved" flag per
+ * {@link ChunkAccess} using the {@link ChunkAccessMixin}, and flush this flag into the
+ * real unsaved flag during "hard" save operations, just before the flag would be
+ * checked. These save operations include:
+ * <ul>
+ * <li>Chunk unloaded.</li>
+ * <li>Game paused (singleplayer).</li>
+ * <li>Save command.</li>
+ * <li>Server stopped.</li>
+ * </ul>
+ * <p>
+ * The flag is set using the injected interface {@link ChunkAccessExt}, via the utility
+ * methods in {@link ChunkUtils}.
+ *
+ * @see ChunkAccessMixin
+ * @see ChunkUtils
+ */
+@Mixin(ChunkMap.class)
+public abstract class ChunkMapMixin extends ChunkStorage {
+    @Shadow private volatile Long2ObjectLinkedOpenHashMap<ChunkHolder> visibleChunkMap;
+
+    public ChunkMapMixin(final Path path, final DataFixer dataFixer, final boolean sync) {
+        super(path, dataFixer, sync);
+    }
+
+    /**
+     * This is for the code-path taken when a chunk is being unloaded.
+     */
+    @Inject(method = "lambda$scheduleUnload$11", at = {@At(value = "INVOKE", target = "Lnet/minecraft/server/level/ChunkMap;save(Lnet/minecraft/world/level/chunk/ChunkAccess;)Z")})
+    private void beforeSaveOnUnload(final ChunkHolder chunkHolder, final CompletableFuture<?> chunkToSave, final long chunkId, final ChunkAccess chunkAccess, final CallbackInfo ci) {
+        if (chunkAccess instanceof ChunkAccessExt ext) {
+            ext.applyAndClearLazyUnsaved();
+        }
+    }
+
+    /**
+     * This is for the code-path taken when saving all chunks upon server shutdown or when
+     * running a save command with the "flush" flag.
+     */
+    @Inject(method = "lambda$saveAllChunks$8", at = {@At(value = "HEAD")})
+    private static void beforeSyncSave(final ChunkAccess chunkAccess, final CallbackInfoReturnable<Boolean> cir) {
+        if (chunkAccess instanceof ChunkAccessExt ext) {
+            ext.applyAndClearLazyUnsaved();
+        }
+    }
+
+    /**
+     * This is for the code-path taken when saving chunk upon pausing the game or when
+     * running a save command without the "flush" flag.
+     */
+    @Inject(method = "saveAllChunks", at = {@At(value = "HEAD")})
+    private void beforeAsyncSave(final boolean sync, final CallbackInfo ci) {
+        // The sync case is handled in beforeSyncSave.
+        if (!sync) {
+            // Need to iterate this ourselves, because I can't find the hook for the save call
+            // inside the foreach in the method. Slightly annoying, but only happens on explicit
+            // save requests, so not too much of a performance worry.
+            visibleChunkMap.values().forEach(holder -> {
+                if (holder.wasAccessibleSinceLastSave()) {
+                    final ChunkAccess chunkToSave = holder.getChunkToSave().getNow(null);
+                    if (chunkToSave instanceof ChunkAccessExt ext) {
+                        ext.applyAndClearLazyUnsaved();
+                    }
+                }
+            });
+        }
+    }
+}

src/main/java/li/cil/oc2/common/mixin/package-info.java

@@ -0,0 +1,7 @@
+@ParametersAreNonnullByDefault
+@MethodsReturnNonnullByDefault
+package li.cil.oc2.common.mixin;
+
+import net.minecraft.MethodsReturnNonnullByDefault;
+
+import javax.annotation.ParametersAreNonnullByDefault;

src/main/java/li/cil/oc2/common/network/Network.java

@@ -2,8 +2,11 @@ package li.cil.oc2.common.network;
 
 import li.cil.oc2.api.API;
 import li.cil.oc2.common.network.message.*;
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.SectionPos;
 import net.minecraft.network.FriendlyByteBuf;
 import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.world.entity.Entity;
 import net.minecraft.world.level.Level;
@@ -86,9 +89,30 @@ public final class Network {
 
     public static <T> void sendToClientsTrackingBlockEntity(final T message, final BlockEntity blockEntity) {
         final Level level = blockEntity.getLevel();
-        if (level != null) {
-            final LevelChunk chunk = level.getChunkAt(blockEntity.getBlockPos());
-            Network.INSTANCE.send(PacketDistributor.TRACKING_CHUNK.with(() -> chunk), message);
+        if (level == null) {
+            return;
+        }
+
+        final MinecraftServer server = level.getServer();
+        if (server == null) {
+            return;
+        }
+
+        if (!server.isSameThread()) {
+            throw new IllegalStateException(
+                "Attempting to send network message to BlockEntity from non-server " +
+                    "thread [" + Thread.currentThread() + "]. This is not supported, " +
+                    "because looking up the chunk from the level is required. " +
+                    "Consider caching the containing chunk and using " +
+                    "sendToClientsTrackingChunk() directly, instead.");
+        }
+
+        final BlockPos blockPos = blockEntity.getBlockPos();
+        final int chunkX = SectionPos.blockToSectionCoord(blockPos.getX());
+        final int chunkZ = SectionPos.blockToSectionCoord(blockPos.getZ());
+        if (level.hasChunk(chunkX, chunkZ)) {
+            final LevelChunk chunk = level.getChunk(chunkX, chunkZ);
+            sendToClientsTrackingChunk(message, chunk);
         }
     }
 

src/main/java/li/cil/oc2/common/util/ChunkAccessExt.java

@@ -0,0 +1,31 @@
+package li.cil.oc2.common.util;
+
+import li.cil.oc2.common.mixin.ChunkMapMixin;
+import net.minecraft.server.level.ChunkMap;
+import net.minecraft.world.level.chunk.ChunkAccess;
+
+/**
+ * Interface injected into the {@link ChunkAccess} class.
+ * <p>
+ * Tracks a "lazy unsaved" flag, which is converted into the regular unsaved flag
+ * before certain manual save operations.
+ *
+ * @see ChunkUtils
+ * @see ChunkMapMixin
+ */
+public interface ChunkAccessExt {
+    /**
+     * Set the lazy unsaved flag for this instance.
+     * <p>
+     * This method is used by the utility methods in {@link ChunkUtils}.
+     */
+    void setLazyUnsaved();
+
+    /**
+     * Set the unsaved flag for this instance, if the lazy unsaved flag is set,
+     * then clears the lazy unsaved flag.
+     * <p>
+     * This method is invoked from mixins injected into the {@link ChunkMap} class.
+     */
+    void applyAndClearLazyUnsaved();
+}

src/main/java/li/cil/oc2/common/util/ChunkUtils.java

@@ -0,0 +1,63 @@
+package li.cil.oc2.common.util;
+
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.SectionPos;
+import net.minecraft.world.level.Level;
+import net.minecraft.world.level.chunk.ChunkAccess;
+
+public final class ChunkUtils {
+    /**
+     * This will mark a chunk unsaved lazily, right before an attempt to save it would be made due
+     * to of these events:
+     * <ul>
+     * <li>Chunk unloaded.</li>
+     * <li>Game paused (singleplayer).</li>
+     * <li>Save command.</li>
+     * <li>Server stopped.</li>
+     * </ul>
+     * <p>
+     * This is intended for things that change every tick, which would lead to saving to NBT every
+     * single tick, when setting {@link net.minecraft.world.level.chunk.ChunkAccess#setUnsaved(boolean)}
+     * directly.
+     * <p>
+     * Instead, this sets a flag on the chunk, which, if true, will cause the chunk to be marked as
+     * unsaved just before this flag is checked, for the events listed above. I.e. for all cases
+     * where an "explicit" save is performed.
+     *
+     * @param chunkAccess the chunk to set the flag for.
+     */
+    public static void setLazyUnsaved(final ChunkAccess chunkAccess) {
+        if (chunkAccess instanceof ChunkAccessExt ext) {
+            ext.setLazyUnsaved();
+        }
+    }
+
+    /**
+     * This will mark a chunk unsaved lazily, right before an attempt to save it would be made due
+     * to of these events:
+     * <ul>
+     * <li>Chunk unloaded.</li>
+     * <li>Game paused (singleplayer).</li>
+     * <li>Save command.</li>
+     * <li>Server stopped.</li>
+     * </ul>
+     * <p>
+     * This is intended for things that change every tick, which would lead to saving to NBT every
+     * single tick, when setting {@link net.minecraft.world.level.chunk.ChunkAccess#setUnsaved(boolean)}
+     * directly.
+     * <p>
+     * Instead, this sets a flag on the chunk, which, if true, will cause the chunk to be marked as
+     * unsaved just before this flag is checked, for the events listed above. I.e. for all cases
+     * where an "explicit" save is performed.
+     *
+     * @param level    the level containing the chunk.
+     * @param blockPos the block position contained in the chunk.
+     */
+    public static void setLazyUnsaved(final Level level, final BlockPos blockPos) {
+        final int chunkX = SectionPos.blockToSectionCoord(blockPos.getX());
+        final int chunkZ = SectionPos.blockToSectionCoord(blockPos.getZ());
+        if (level.hasChunk(chunkX, chunkZ)) {
+            setLazyUnsaved(level.getChunk(chunkX, chunkZ));
+        }
+    }
+}

src/main/resources/oc2.mixins.json

@@ -0,0 +1,14 @@
+{
+    "minVersion": "0.8",
+    "compatibilityLevel": "JAVA_17",
+    "required": true,
+    "package": "li.cil.oc2.common.mixin",
+    "refmap": "oc2.refmap.json",
+    "mixins": [
+        "ChunkAccessMixin",
+        "ChunkMapMixin"
+    ],
+    "injectors": {
+        "defaultRequire": 1
+    }
+}