From aa3e3b9b14ee7c837444c74d0c4519cb934fa75b Mon Sep 17 00:00:00 2001
From: Spottedleaf <Spottedleaf@users.noreply.github.com>
Date: Tue, 25 Jul 2023 07:11:58 -0700
Subject: [PATCH] Prevent entity removals if the entity slices is receiving
 status updates

If an entity is removed while updating an entity slice, then the
iteration over the entity slice's entities could throw a cryptic
exception. Instead, it is better to prevent the entity removal
with a useful log message.

Fixes https://github.com/PaperMC/Paper/issues/9464
---
 patches/server/Rewrite-chunk-system.patch | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/patches/server/Rewrite-chunk-system.patch b/patches/server/Rewrite-chunk-system.patch
index cab068f113..7760387c9b 100644
--- a/patches/server/Rewrite-chunk-system.patch
+++ b/patches/server/Rewrite-chunk-system.patch
@@ -4144,6 +4144,17 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +        return true;
 +    }
 +
++    public boolean canRemoveEntity(final Entity entity) {
++        if (entity.updatingSectionStatus) {
++            return false;
++        }
++
++        final int sectionX = entity.sectionX;
++        final int sectionZ = entity.sectionZ;
++        final ChunkEntitySlices slices = this.getChunk(sectionX, sectionZ);
++        return slices == null || !slices.isPreventingStatusUpdates();
++    }
++
 +    private void removeEntity(final Entity entity) {
 +        final int sectionX = entity.sectionX;
 +        final int sectionY = entity.sectionY;
@@ -4157,6 +4168,9 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +        if (slices == null) {
 +            LOGGER.warn("Cannot remove entity " + entity + " from null entity slices (" + sectionX + "," + sectionZ + ")");
 +        } else {
++            if (slices.isPreventingStatusUpdates()) {
++                throw new IllegalStateException("Attempting to remove entity " + entity + " from entity slices (" + sectionX + "," + sectionZ + ") that is receiving status updates");
++            }
 +            if (!slices.removeEntity(entity, sectionY)) {
 +                LOGGER.warn("Failed to remove entity " + entity + " from entity slices (" + sectionX + "," + sectionZ + ")");
 +            }
@@ -16307,6 +16321,10 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
 +        return ret;
 +    }
 +
++    public boolean isPreventingStatusUpdates() {
++        return this.preventStatusUpdates;
++    }
++
 +    public void stopPreventingStatusUpdates(final boolean prev) {
 +        this.preventStatusUpdates = prev;
 +    }
@@ -20748,8 +20766,8 @@ index 0000000000000000000000000000000000000000..00000000000000000000000000000000
      public final void setRemoved(Entity.RemovalReason reason) {
 +        // Paper start - rewrite chunk system
 +        io.papermc.paper.util.TickThread.ensureTickThread(this, "Cannot remove entity off-main");
-+        if (this.updatingSectionStatus) {
-+            LOGGER.warn("Entity " + this + " is currently prevented from being added/removed to world since it is processing section status updates", new Throwable());
++        if (!((ServerLevel)this.level).getEntityLookup().canRemoveEntity(this)) {
++            LOGGER.warn("Entity " + this + " is currently prevented from being removed from the world since it is processing section status updates", new Throwable());
 +            return;
 +        }
 +        // Paper end - rewrite chunk system