Reschedule block entities when chunks are loaded

Minecraft sometimes keeps chunks in-memory, but not actively loaded. If
we schedule a block entity to be ticked and that chunk is is then
transitioned to this partially-loaded state, then the block entity is
never actually ticked.

This is most visible with monitors. When a monitor's contents changes,
if the monitor is not already marked as changed, we set it as changed
and schedule a tick (see ServerMonitor). However, if the tick is
dropped, we don't clear the changed flag, meaning subsequent changes
don't requeue the monitor to be ticked, and so the monitor is never
updated.

We fix this by maintaining a list of block entities whose tick was
dropped. If these block entities (or rather their owning chunk) is ever
re-loaded, then we reschedule them to be ticked.

An alternative approach here would be to add the scheduled tick directly
to the LevelChunk. However, getting hold of the LevelChunk for unloaded
blocks is quiet nasty, so I think best avoided.

Fixes #1146. Fixes #1560 - I believe the second one is a duplicate, and
I noticed too late :D.

projects/common/src/main/java/dan200/computercraft/shared/CommonHooks.java

@@ -18,6 +18,7 @@
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.dedicated.DedicatedServer;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.server.packs.resources.PreparableReloadListener;
 import net.minecraft.world.entity.Entity;
@@ -72,10 +73,19 @@ private static void resetState() {
         NetworkUtils.reset();
     }
 
+    public static void onServerChunkUnload(LevelChunk chunk) {
+        if (!(chunk.getLevel() instanceof ServerLevel)) throw new IllegalArgumentException("Not a server chunk.");
+        TickScheduler.onChunkUnload(chunk);
+    }
+
     public static void onChunkWatch(LevelChunk chunk, ServerPlayer player) {
         MonitorWatcher.onWatch(chunk, player);
     }
 
+    public static void onChunkTicketLevelChanged(ServerLevel level, long chunkPos, int oldLevel, int newLevel) {
+        TickScheduler.onChunkTicketChanged(level, chunkPos, oldLevel, newLevel);
+    }
+
     public static final ResourceLocation TREASURE_DISK_LOOT = new ResourceLocation(ComputerCraftAPI.MOD_ID, "treasure_disk");
 
     private static final Set<ResourceLocation> TREASURE_DISK_LOOT_TABLES = Set.of(

projects/common/src/main/java/dan200/computercraft/shared/peripheral/modem/wired/CableBlockEntity.java

@@ -265,7 +265,7 @@ void connectionsChanged() {
                 // If we can connect to it then do so
                 this.node.connectTo(node);
             } else {
-                // Otherwise break the connectoin.
+                // Otherwise break the connection.
                 this.node.disconnectFrom(node);
             }
         }

projects/common/src/main/java/dan200/computercraft/shared/util/TickScheduler.java

@@ -5,13 +5,20 @@
 package dan200.computercraft.shared.util;
 
 import net.minecraft.core.BlockPos;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.server.level.ChunkLevel;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.level.ChunkPos;
+import net.minecraft.world.level.Level;
 import net.minecraft.world.level.LevelAccessor;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.entity.BlockEntity;
+import net.minecraft.world.level.chunk.ChunkStatus;
+import net.minecraft.world.level.chunk.LevelChunk;
 
-import java.util.Queue;
+import java.util.*;
 import java.util.concurrent.ConcurrentLinkedDeque;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 
 /**
  * A thread-safe version of {@link LevelAccessor#scheduleTick(BlockPos, Block, int)}.
@@ -22,26 +29,92 @@ public final class TickScheduler {
     private TickScheduler() {
     }
 
+    /**
+     * The list of block entities to tick.
+     */
     private static final Queue<Token> toTick = new ConcurrentLinkedDeque<>();
 
+    /**
+     * Block entities which we want to tick, but whose chunks not currently loaded.
+     * <p>
+     * Minecraft sometimes keeps chunks in-memory, but not actively loaded. If such a block entity is in the
+     * {@link #toTick} queue, we'll see that it's not loaded and so have to skip scheduling a tick.
+     * <p>
+     * However, if the block entity is ever loaded again, we need to tick it. Unfortunately, block entities in this
+     * state are not notified in any way (for instance, {@link BlockEntity#setRemoved()} or
+     * {@link BlockEntity#clearRemoved()} are not called), and so there's no way to easily reschedule them for ticking.
+     * <p>
+     * Instead, for each chunk we keep a list of all block entities whose tick we skipped. If a chunk is loaded,
+     * {@linkplain #onChunkTicketChanged(ServerLevel, long, int, int) we requeue all skipped ticks}.
+     */
+    private static final Map<ChunkReference, List<Token>> delayed = new HashMap<>();
+
+    /**
+     * Schedule a block entity to be ticked.
+     *
+     * @param token The token whose block entity should be ticked.
+     */
     public static void schedule(Token token) {
         var world = token.owner.getLevel();
-        if (world != null && !world.isClientSide && !token.scheduled.getAndSet(true)) toTick.add(token);
+        if (world != null && !world.isClientSide && Token.STATE.compareAndSet(token, State.IDLE, State.SCHEDULED)) {
+            toTick.add(token);
+        }
+    }
+
+    public static void onChunkTicketChanged(ServerLevel level, long chunkPos, int oldLevel, int newLevel) {
+        boolean oldLoaded = isLoaded(oldLevel), newLoaded = isLoaded(newLevel);
+        if (!oldLoaded && newLoaded) {
+            // If our chunk is becoming active, requeue all pending tokens.
+            var delayedTokens = delayed.remove(new ChunkReference(level.dimension(), chunkPos));
+            if (delayedTokens == null) return;
+
+            for (var token : delayedTokens) {
+                if (token.owner.isRemoved()) {
+                    Token.STATE.set(token, State.IDLE);
+                } else {
+                    Token.STATE.set(token, State.SCHEDULED);
+                    toTick.add(token);
+                }
+            }
+        }
+    }
+
+    public static void onChunkUnload(LevelChunk chunk) {
+        // If our chunk is fully unloaded, all block entities are about to be removed - we need to dequeue any delayed
+        // tokens from the queue.
+        var delayedTokens = delayed.remove(new ChunkReference(chunk.getLevel().dimension(), chunk.getPos().toLong()));
+        if (delayedTokens == null) return;
+
+        for (var token : delayedTokens) Token.STATE.set(token, State.IDLE);
     }
 
     public static void tick() {
         Token token;
-        while ((token = toTick.poll()) != null) {
-            token.scheduled.set(false);
-            var blockEntity = token.owner;
-            if (blockEntity.isRemoved()) continue;
+        while ((token = toTick.poll()) != null) Token.STATE.set(token, tickToken(token));
+    }
+
+    private static State tickToken(Token token) {
+        var blockEntity = token.owner;
+
+        // If the block entity has been removed, then remove it from the queue.
+        if (blockEntity.isRemoved()) return State.IDLE;
 
-            var world = blockEntity.getLevel();
-            var pos = blockEntity.getBlockPos();
+        var level = Objects.requireNonNull(blockEntity.getLevel(), "Block entity level cannot become null");
+        var pos = blockEntity.getBlockPos();
 
-            if (world != null && world.isLoaded(pos) && world.getBlockEntity(pos) == blockEntity) {
-                world.scheduleTick(pos, blockEntity.getBlockState().getBlock(), 0);
+        if (!level.isLoaded(pos)) {
+            // The chunk is not properly loaded, as it to our delayed set.
+            delayed.computeIfAbsent(new ChunkReference(level.dimension(), ChunkPos.asLong(pos)), x -> new ArrayList<>()).add(token);
+            return State.UNLOADED;
+        } else {
+            // This should be impossible: either the block entity is at the above position, or it has been removed.
+            if (level.getBlockEntity(pos) != blockEntity) {
+                throw new IllegalStateException("Expected " + blockEntity + " at " + pos);
             }
+
+            // Otherwise schedule a tick and remove it from the queue.
+            level.scheduleTick(pos, blockEntity.getBlockState().getBlock(), 0);
+            return State.IDLE;
         }
     }
 
@@ -52,11 +125,51 @@ public static void tick() {
      * As such, it should be unique per {@link BlockEntity} instance to avoid it being queued multiple times.
      */
     public static class Token {
+        static final AtomicReferenceFieldUpdater<Token, State> STATE = AtomicReferenceFieldUpdater.newUpdater(Token.class, State.class, "$state");
+
         final BlockEntity owner;
-        final AtomicBoolean scheduled = new AtomicBoolean();
+
+        /**
+         * The current state of this token.
+         */
+        private volatile State $state = State.IDLE;
 
         public Token(BlockEntity owner) {
             this.owner = owner;
         }
     }
+
+    /**
+     * The possible states a {@link Token} can be in.
+     * <p>
+     * This effectively stores which (if any) queue the token is currently in, allowing us to skip scheduling if the
+     * token is already enqueued.
+     */
+    private enum State {
+        /**
+         * The token is not on any queues.
+         */
+        IDLE,
+
+        /**
+         * The token is on the {@link #toTick} queue.
+         */
+        SCHEDULED,
+
+        /**
+         * The token is on the {@link #delayed} queue.
+         */
+        UNLOADED,
+    }
+
+    private record ChunkReference(ResourceKey<Level> level, Long position) {
+        @Override
+        public String toString() {
+            return "ChunkReference(" + level + " at " + new ChunkPos(position) + ")";
+        }
+    }
+
+    private static boolean isLoaded(int level) {
+        return level <= ChunkLevel.byStatus(ChunkStatus.FULL);
+    }
 }

projects/fabric/src/main/java/dan200/computercraft/mixin/ChunkMapMixin.java

@@ -6,20 +6,37 @@
 
 import dan200.computercraft.shared.CommonHooks;
 import net.minecraft.network.protocol.game.ClientboundLevelChunkWithLightPacket;
+import net.minecraft.server.level.ChunkHolder;
 import net.minecraft.server.level.ChunkMap;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.world.level.chunk.LevelChunk;
 import org.apache.commons.lang3.mutable.MutableObject;
+import org.spongepowered.asm.mixin.Final;
 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 javax.annotation.Nullable;
 
 @Mixin(ChunkMap.class)
 class ChunkMapMixin {
+    @Final
+    @Shadow
+    ServerLevel level;
+
     @Inject(method = "playerLoadedChunk", at = @At("TAIL"))
     @SuppressWarnings("UnusedMethod")
     private void onPlayerLoadedChunk(ServerPlayer player, MutableObject<ClientboundLevelChunkWithLightPacket> packetCache, LevelChunk chunk, CallbackInfo callback) {
         CommonHooks.onChunkWatch(chunk, player);
     }
+
+    @Inject(method = "updateChunkScheduling", at = @At("HEAD"))
+    @SuppressWarnings("UnusedMethod")
+    private void onUpdateChunkScheduling(long chunkPos, int newLevel, @Nullable ChunkHolder holder, int oldLevel, CallbackInfoReturnable<ChunkHolder> callback) {
+        CommonHooks.onChunkTicketLevelChanged(level, chunkPos, oldLevel, newLevel);
+    }
 }

projects/fabric/src/main/java/dan200/computercraft/shared/ComputerCraft.java

@@ -24,6 +24,7 @@
 import dan200.computercraft.shared.platform.FabricConfigFile;
 import dan200.computercraft.shared.platform.FabricMessageType;
 import net.fabricmc.fabric.api.command.v2.CommandRegistrationCallback;
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerChunkEvents;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerTickEvents;
 import net.fabricmc.fabric.api.event.player.PlayerBlockBreakEvents;
@@ -96,6 +97,7 @@ public static void init() {
 
         ServerTickEvents.START_SERVER_TICK.register(CommonHooks::onServerTickStart);
         ServerTickEvents.START_SERVER_TICK.register(s -> CommonHooks.onServerTickEnd());
+        ServerChunkEvents.CHUNK_UNLOAD.register((l, c) -> CommonHooks.onServerChunkUnload(c));
 
         PlayerBlockBreakEvents.BEFORE.register(FabricCommonHooks::onBlockDestroy);
         UseBlockCallback.EVENT.register(FabricCommonHooks::useOnBlock);

projects/forge/src/main/java/dan200/computercraft/shared/ForgeCommonHooks.java

@@ -22,11 +22,15 @@
 import dan200.computercraft.shared.util.CapabilityProvider;
 import dan200.computercraft.shared.util.SidedCapabilityProvider;
 import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.level.block.entity.BlockEntity;
 import net.minecraft.world.level.block.entity.CommandBlockEntity;
+import net.minecraft.world.level.chunk.LevelChunk;
 import net.minecraftforge.event.*;
 import net.minecraftforge.event.entity.EntityJoinLevelEvent;
 import net.minecraftforge.event.entity.living.LivingDropsEvent;
+import net.minecraftforge.event.level.ChunkEvent;
+import net.minecraftforge.event.level.ChunkTicketLevelUpdatedEvent;
 import net.minecraftforge.event.level.ChunkWatchEvent;
 import net.minecraftforge.event.server.ServerStartingEvent;
 import net.minecraftforge.event.server.ServerStoppedEvent;
@@ -67,11 +71,23 @@ public static void onRegisterCommand(RegisterCommandsEvent event) {
         CommandComputerCraft.register(event.getDispatcher());
     }
 
+    @SubscribeEvent
+    public static void onChunkUnload(ChunkEvent.Unload event) {
+        if (event.getLevel() instanceof ServerLevel && event.getChunk() instanceof LevelChunk chunk) {
+            CommonHooks.onServerChunkUnload(chunk);
+        }
+    }
+
     @SubscribeEvent
     public static void onChunkWatch(ChunkWatchEvent.Watch event) {
         CommonHooks.onChunkWatch(event.getChunk(), event.getPlayer());
     }
 
+    @SubscribeEvent
+    public static void onChunkTicketLevelChanged(ChunkTicketLevelUpdatedEvent event) {
+        CommonHooks.onChunkTicketLevelChanged(event.getLevel(), event.getChunkPos(), event.getOldTicketLevel(), event.getNewTicketLevel());
+    }
+
     @SubscribeEvent
     public static void onAddReloadListeners(AddReloadListenerEvent event) {
         CommonHooks.onDatapackReload((id, listener) -> event.addListener(listener));