From 196bf7034d4190cc62a8e89304965262681f559d Mon Sep 17 00:00:00 2001
From: jebibot <83044352+jebibot@users.noreply.github.com>
Date: Sat, 5 Aug 2023 12:52:23 +0800
Subject: [PATCH] Fix CME when plugin load chunks during chunk load (#1056)

---
 .../server/level/DistanceManagerMixin.java    | 29 +++++++++++++++----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/arclight-common/src/main/java/io/izzel/arclight/common/mixin/core/server/level/DistanceManagerMixin.java b/arclight-common/src/main/java/io/izzel/arclight/common/mixin/core/server/level/DistanceManagerMixin.java
index 37ab1a0a..f9b138b4 100644
--- a/arclight-common/src/main/java/io/izzel/arclight/common/mixin/core/server/level/DistanceManagerMixin.java
+++ b/arclight-common/src/main/java/io/izzel/arclight/common/mixin/core/server/level/DistanceManagerMixin.java
@@ -5,11 +5,7 @@ import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.ObjectSet;
 import net.minecraft.core.SectionPos;
-import net.minecraft.server.level.DistanceManager;
-import net.minecraft.server.level.ServerPlayer;
-import net.minecraft.server.level.Ticket;
-import net.minecraft.server.level.TicketType;
-import net.minecraft.server.level.TickingTracker;
+import net.minecraft.server.level.*;
 import net.minecraft.util.SortedArraySet;
 import net.minecraft.world.level.ChunkPos;
 import org.spongepowered.asm.mixin.Final;
@@ -18,10 +14,13 @@ import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.gen.Invoker;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
 
 import java.util.Iterator;
+import java.util.Set;
+import java.util.function.Consumer;
 
 @Mixin(DistanceManager.class)
 public abstract class DistanceManagerMixin implements TicketManagerBridge {
@@ -44,6 +43,26 @@ public abstract class DistanceManagerMixin implements TicketManagerBridge {
         }
     }
 
+    @Redirect(method = "runAllUpdates", at = @At(value = "INVOKE", remap = false, target = "Ljava/util/Set;forEach(Ljava/util/function/Consumer;)V"))
+    private void arclight$safeIter(Set<ChunkHolder> instance, Consumer<ChunkHolder> consumer) {
+        // Iterate pending chunk updates with protection against concurrent modification exceptions
+        var iter = instance.iterator();
+        var expectedSize = instance.size();
+        do {
+            var chunkHolder = iter.next();
+            iter.remove();
+            expectedSize--;
+
+            consumer.accept(chunkHolder);
+
+            // Reset iterator if set was modified using add()
+            if (instance.size() != expectedSize) {
+                expectedSize = instance.size();
+                iter = instance.iterator();
+            }
+        } while (iter.hasNext());
+    }
+
     public <T> boolean addRegionTicketAtDistance(TicketType<T> type, ChunkPos pos, int level, T value) {
         var ticket = new Ticket<>(type, 33 - level, value);
         var ret = this.addTicket(pos.toLong(), ticket);