Add debug feature to catch illegal block entity map mutation

Author: embeddedt

common/src/main/java/org/embeddedt/modernfix/common/mixin/feature/blockentity_incorrect_thread/ChunkAccessMixin.java

@@ -0,0 +1,34 @@
+package org.embeddedt.modernfix.common.mixin.feature.blockentity_incorrect_thread;
+
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.Registry;
+import net.minecraft.world.level.ChunkPos;
+import net.minecraft.world.level.Level;
+import net.minecraft.world.level.LevelHeightAccessor;
+import net.minecraft.world.level.block.entity.BlockEntity;
+import net.minecraft.world.level.chunk.ChunkAccess;
+import net.minecraft.world.level.chunk.LevelChunkSection;
+import net.minecraft.world.level.chunk.UpgradeData;
+import net.minecraft.world.level.levelgen.blending.BlendingData;
+import org.embeddedt.modernfix.util.ConcurrencySanitizingMap;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Mutable;
+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 java.util.Map;
+
+@Mixin(ChunkAccess.class)
+public class ChunkAccessMixin {
+    @Shadow @Final @Mutable protected Map<BlockPos, BlockEntity> blockEntities;
+
+    @Inject(method = "<init>", at = @At("RETURN"))
+    private void wrapInConcurrencyDetector(ChunkPos chunkPos, UpgradeData upgradeData, LevelHeightAccessor levelHeightAccessor, Registry biomeRegistry, long inhabitedTime, LevelChunkSection[] sections, BlendingData blendingData, CallbackInfo ci) {
+        if (levelHeightAccessor instanceof Level level) {
+            this.blockEntities = new ConcurrencySanitizingMap<>(this.blockEntities, ((LevelThreadAccessor)level).getThread());
+        }
+    }
+}

common/src/main/java/org/embeddedt/modernfix/common/mixin/feature/blockentity_incorrect_thread/LevelThreadAccessor.java

@@ -0,0 +1,11 @@
+package org.embeddedt.modernfix.common.mixin.feature.blockentity_incorrect_thread;
+
+import net.minecraft.world.level.Level;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(Level.class)
+public interface LevelThreadAccessor {
+    @Accessor
+    Thread getThread();
+}

common/src/main/java/org/embeddedt/modernfix/core/config/ModernFixEarlyConfig.java

@@ -168,6 +168,7 @@ public DefaultSettingMapBuilder put(String key, Boolean value) {
             .put("mixin.perf.worldgen_allocation", false) // experimental
             .put("mixin.feature.cause_lag_by_disabling_threads", false)
             .put("mixin.bugfix.missing_block_entities", false)
+            .put("mixin.feature.blockentity_incorrect_thread", false)
             .put("mixin.perf.clear_mixin_classinfo", false)
             .put("mixin.perf.deduplicate_climate_parameters", false)
             .put("mixin.bugfix.packet_leak", false)

common/src/main/java/org/embeddedt/modernfix/util/ConcurrencySanitizingMap.java

@@ -0,0 +1,100 @@
+package org.embeddedt.modernfix.util;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * A map wrapper that throws if the map is ever accessed on the wrong thread.
+ */
+public class ConcurrencySanitizingMap<K, V> implements Map<K, V> {
+    private final Map<K, V> map;
+    private final Thread owner;
+
+    public ConcurrencySanitizingMap(Map<K, V> map, Thread owner) {
+        this.map = map;
+        this.owner = owner;
+    }
+
+    private void checkThread() {
+        var current = Thread.currentThread();
+        if (current != owner) {
+            throw new IllegalStateException("Map is being accessed on thread " + current + " while owned by " + owner);
+        }
+    }
+
+    @Override
+    public int size() {
+        checkThread();
+        return map.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        checkThread();
+        return map.isEmpty();
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        checkThread();
+        return map.containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        checkThread();
+        return map.containsValue(value);
+    }
+
+    @Override
+    public V get(Object key) {
+        checkThread();
+        return map.get(key);
+    }
+
+    @Override
+    public @Nullable V put(K key, V value) {
+        checkThread();
+        return map.put(key, value);
+    }
+
+    @Override
+    public V remove(Object key) {
+        checkThread();
+        return map.remove(key);
+    }
+
+    @Override
+    public void putAll(@NotNull Map<? extends K, ? extends V> m) {
+        checkThread();
+        map.putAll(m);
+    }
+
+    @Override
+    public void clear() {
+        checkThread();
+        map.clear();
+    }
+
+    @Override
+    public @NotNull Set<K> keySet() {
+        checkThread();
+        return map.keySet();
+    }
+
+    @Override
+    public @NotNull Collection<V> values() {
+        checkThread();
+        return map.values();
+    }
+
+    @Override
+    public @NotNull Set<Entry<K, V>> entrySet() {
+        checkThread();
+        return map.entrySet();
+    }
+}