SPIGOT-6547: Chunk#getEntities() doesn't return all entities immediately after chunk load

nms-patches/net/minecraft/world/level/chunk/storage/EntityStorage.patch

@@ -0,0 +1,22 @@
+--- a/net/minecraft/world/level/chunk/storage/EntityStorage.java
++++ b/net/minecraft/world/level/chunk/storage/EntityStorage.java
+@@ -33,7 +33,7 @@
+     private static final Logger LOGGER = LogManager.getLogger();
+     private static final String ENTITIES_TAG = "Entities";
+     private static final String POSITION_TAG = "Position";
+-    private final WorldServer level;
++    public final WorldServer level; // PAIL private -> public
+     private final IOWorker worker;
+     private final LongSet emptyChunks = new LongOpenHashSet();
+     private final ThreadedMailbox<Runnable> entityDeserializerQueue;
+@@ -51,8 +51,8 @@
+         if (this.emptyChunks.contains(chunkcoordintpair.pair())) {
+             return CompletableFuture.completedFuture(b(chunkcoordintpair));
+         } else {
+-            CompletableFuture completablefuture = this.worker.b(chunkcoordintpair);
+-            Function function = (nbttagcompound) -> {
++            CompletableFuture<NBTTagCompound> completablefuture = this.worker.b(chunkcoordintpair); // CraftBukkit - decompile error
++            Function<NBTTagCompound, ChunkEntities<Entity>> function = (nbttagcompound) -> { // CraftBukkit - decompile error
+                 if (nbttagcompound == null) {
+                     this.emptyChunks.add(chunkcoordintpair.pair());
+                     return b(chunkcoordintpair);

nms-patches/net/minecraft/world/level/entity/PersistentEntitySectionManager.patch

@@ -1,6 +1,31 @@
 --- a/net/minecraft/world/level/entity/PersistentEntitySectionManager.java
 +++ b/net/minecraft/world/level/entity/PersistentEntitySectionManager.java
-@@ -82,7 +82,7 @@
+@@ -32,6 +32,11 @@
+ import org.apache.logging.log4j.LogManager;
+ import org.apache.logging.log4j.Logger;
+ 
++// CraftBukkit start
++import net.minecraft.world.level.chunk.storage.EntityStorage;
++import org.bukkit.craftbukkit.event.CraftEventFactory;
++// CraftBukkit end
++
+ public class PersistentEntitySectionManager<T extends EntityAccess> implements AutoCloseable {
+ 
+     static final Logger LOGGER = LogManager.getLogger();
+@@ -55,6 +60,12 @@
+         this.entityGetter = new LevelEntityGetterAdapter<>(this.visibleEntityStorage, this.sectionStorage);
+     }
+ 
++    // CraftBukkit start - add method to get all entities in chunk
++    public List<Entity> getEntities(ChunkCoordIntPair chunkCoordIntPair) {
++        return sectionStorage.b(chunkCoordIntPair.pair()).flatMap(EntitySection::b).map(entiy -> (Entity) entiy).collect(Collectors.toList());
++    }
++    // CraftBukkit end
++
+     void a(long i, EntitySection<T> entitysection) {
+         if (entitysection.a()) {
+             this.sectionStorage.e(i);
+@@ -82,7 +93,7 @@
              long i = SectionPosition.c(t0.getChunkCoordinates());
              EntitySection<T> entitysection = this.sectionStorage.c(i);
  
@@ -9,7 +34,54 @@
              t0.a(new PersistentEntitySectionManager.a(t0, i, entitysection));
              if (!flag) {
                  this.callbacks.f(t0);
-@@ -254,13 +254,13 @@
+@@ -186,7 +197,7 @@
+         });
+     }
+ 
+-    private void b(long i) {
++    public void b(long i) { // PAIL private -> public, rename scheduleEntityLoading
+         PersistentEntitySectionManager.b persistententitysectionmanager_b = (PersistentEntitySectionManager.b) this.chunkLoadStatuses.get(i);
+ 
+         if (persistententitysectionmanager_b == PersistentEntitySectionManager.b.FRESH) {
+@@ -196,6 +207,12 @@
+     }
+ 
+     private boolean a(long i, Consumer<T> consumer) {
++        // CraftBukkit start - add boolean for event call
++        return a(i, consumer, false);
++    }
++
++    private boolean a(long i, Consumer<T> consumer, boolean callEvent) {
++        // CraftBukkit end
+         PersistentEntitySectionManager.b persistententitysectionmanager_b = (PersistentEntitySectionManager.b) this.chunkLoadStatuses.get(i);
+ 
+         if (persistententitysectionmanager_b == PersistentEntitySectionManager.b.PENDING) {
+@@ -207,6 +224,7 @@
+ 
+             if (list.isEmpty()) {
+                 if (persistententitysectionmanager_b == PersistentEntitySectionManager.b.LOADED) {
++                    if (callEvent) CraftEventFactory.callEntitiesUnloadEvent(((EntityStorage) permanentStorage).level, new ChunkCoordIntPair(i), ImmutableList.of()); // CraftBukkit
+                     this.permanentStorage.a(new ChunkEntities<>(new ChunkCoordIntPair(i), ImmutableList.of()));
+                 }
+ 
+@@ -215,6 +233,7 @@
+                 this.c(i);
+                 return false;
+             } else {
++                if (callEvent) CraftEventFactory.callEntitiesUnloadEvent(((EntityStorage) permanentStorage).level, new ChunkCoordIntPair(i), list.stream().map(entity -> (Entity) entity).collect(Collectors.toList())); // CraftBukkit
+                 this.permanentStorage.a(new ChunkEntities<>(new ChunkCoordIntPair(i), list));
+                 list.forEach(consumer);
+                 return true;
+@@ -238,7 +257,7 @@
+     private boolean d(long i) {
+         boolean flag = this.a(i, (entityaccess) -> {
+             entityaccess.cD().forEach(this::g);
+-        });
++        }, true); // CraftBukkit - add boolean for event call
+ 
+         if (!flag) {
+             return false;
+@@ -254,19 +273,23 @@
      }
  
      private void f() {
@@ -25,7 +97,17 @@
  
          while ((chunkentities = (ChunkEntities) this.loadingInbox.poll()) != null) {
              chunkentities.b().forEach((entityaccess) -> {
-@@ -292,7 +292,7 @@
+                 this.a(entityaccess, true);
+             });
+             this.chunkLoadStatuses.put(chunkentities.a().pair(), PersistentEntitySectionManager.b.LOADED);
++            // CraftBukkit start - call entity load event
++            List<Entity> entities = getEntities(chunkentities.a()); // PAIL rename getChunkPos
++            CraftEventFactory.callEntitiesLoadEvent(((EntityStorage) permanentStorage).level, chunkentities.a(), entities);
++            // CraftBukkit end
+         }
+ 
+     }
+@@ -292,7 +315,7 @@
      }
  
      public void b() {
@@ -34,7 +116,7 @@
              boolean flag = this.chunkVisibility.get(i) == Visibility.HIDDEN;
  
              if (flag) {
-@@ -311,7 +311,7 @@
+@@ -311,7 +334,7 @@
          while (!longset.isEmpty()) {
              this.permanentStorage.a(false);
              this.g();
@@ -43,7 +125,7 @@
                  boolean flag = this.chunkVisibility.get(i) == Visibility.HIDDEN;
  
                  return flag ? this.d(i) : this.a(i, (entityaccess) -> {
-@@ -323,7 +323,15 @@
+@@ -323,7 +346,15 @@
      }
  
      public void close() throws IOException {
@@ -60,7 +142,7 @@
          this.permanentStorage.close();
      }
  
-@@ -350,7 +358,7 @@
+@@ -350,7 +381,7 @@
      public void a(Writer writer) throws IOException {
          CSVWriter csvwriter = CSVWriter.a().a("x").a("y").a("z").a("visibility").a("load_status").a("entity_count").a(writer);
  
@@ -69,7 +151,7 @@
              PersistentEntitySectionManager.b persistententitysectionmanager_b = (PersistentEntitySectionManager.b) this.chunkLoadStatuses.get(i);
  
              this.sectionStorage.a(i).forEach((j) -> {
-@@ -389,7 +397,7 @@
+@@ -389,7 +420,7 @@
          private EntitySection<T> currentSection;
  
          a(EntityAccess entityaccess, long i, EntitySection entitysection) {
@@ -78,7 +160,7 @@
              this.currentSectionKey = i;
              this.currentSection = entitysection;
          }
-@@ -409,7 +417,7 @@
+@@ -409,7 +440,7 @@
                  PersistentEntitySectionManager.this.a(this.currentSectionKey, this.currentSection);
                  EntitySection<T> entitysection = PersistentEntitySectionManager.this.sectionStorage.c(i);
  

src/main/java/org/bukkit/craftbukkit/CraftChunk.java

@@ -5,6 +5,7 @@ import com.google.common.base.Predicates;
 import java.lang.ref.WeakReference;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Objects;
 import java.util.function.Predicate;
 import net.minecraft.core.BlockPosition;
 import net.minecraft.core.IRegistry;
@@ -23,12 +24,12 @@ import net.minecraft.world.level.chunk.ChunkStatus;
 import net.minecraft.world.level.chunk.DataPaletteBlock;
 import net.minecraft.world.level.chunk.IChunkAccess;
 import net.minecraft.world.level.chunk.NibbleArray;
+import net.minecraft.world.level.entity.PersistentEntitySectionManager;
 import net.minecraft.world.level.levelgen.HeightMap;
 import net.minecraft.world.level.levelgen.SeededRandom;
 import net.minecraft.world.level.lighting.LightEngine;
 import org.bukkit.Chunk;
 import org.bukkit.ChunkSnapshot;
-import org.bukkit.Location;
 import org.bukkit.World;
 import org.bukkit.block.Block;
 import org.bukkit.block.BlockState;
@@ -56,6 +57,13 @@ public class CraftChunk implements Chunk {
         z = getHandle().getPos().z;
     }
 
+    public CraftChunk(WorldServer worldServer, int x, int z) {
+        this.weakChunk = new WeakReference<>(null);
+        this.worldServer = worldServer;
+        this.x = x;
+        this.z = z;
+    }
+
     @Override
     public World getWorld() {
         return worldServer.getWorld();
@@ -109,11 +117,35 @@ public class CraftChunk implements Chunk {
             getWorld().getChunkAt(x, z); // Transient load for this tick
         }
 
-        Location location = new Location(null, 0, 0, 0);
+        PersistentEntitySectionManager<net.minecraft.world.entity.Entity> entityManager = getCraftWorld().getHandle().entityManager;
-        return getWorld().getEntities().stream().filter((entity) -> {
+        long pair = ChunkCoordIntPair.pair(x, z);
-            entity.getLocation(location);
+
-            return location.getBlockX() >> 4 == this.x && location.getBlockZ() >> 4 == this.z;
+        if (entityManager.a(pair)) { // PAIL rename isEntitiesLoaded
-        }).toArray(Entity[]::new);
+            return entityManager.getEntities(new ChunkCoordIntPair(x, z)).stream()
+                    .map(net.minecraft.world.entity.Entity::getBukkitEntity)
+                    .filter(Objects::nonNull).toArray(Entity[]::new);
+        }
+
+        entityManager.b(pair); // Start entity loading
+
+        // now we wait until the entities are loaded,
+        // the converting from NBT to entity object is done on the main Thread which is why we wait
+        getCraftWorld().getHandle().getMinecraftServer().awaitTasks(() -> {
+            boolean status = entityManager.a(pair);
+            // only execute inbox if our entities are not present
+            if (status) {
+                return true;
+            }
+            // tick loading inbox, which loads the created entities to the world
+            // (if present)
+            entityManager.tick();
+            // check if our entities are loaded
+            return entityManager.a(pair);
+        });
+
+        return entityManager.getEntities(new ChunkCoordIntPair(x, z)).stream()
+                .map(net.minecraft.world.entity.Entity::getBukkitEntity)
+                .filter(Objects::nonNull).toArray(Entity[]::new);
     }
 
     @Override

src/main/java/org/bukkit/craftbukkit/event/CraftEventFactory.java

@@ -56,6 +56,7 @@ import net.minecraft.world.inventory.ContainerMerchant;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.Items;
 import net.minecraft.world.item.context.ItemActionContext;
+import net.minecraft.world.level.ChunkCoordIntPair;
 import net.minecraft.world.level.Explosion;
 import net.minecraft.world.level.GeneratorAccess;
 import net.minecraft.world.level.World;
@@ -75,6 +76,7 @@ import org.bukkit.Statistic.Type;
 import org.bukkit.block.Block;
 import org.bukkit.block.BlockFace;
 import org.bukkit.block.BlockState;
+import org.bukkit.craftbukkit.CraftChunk;
 import org.bukkit.craftbukkit.CraftLootTable;
 import org.bukkit.craftbukkit.CraftRaid;
 import org.bukkit.craftbukkit.CraftServer;
@@ -219,6 +221,8 @@ import org.bukkit.event.raid.RaidTriggerEvent;
 import org.bukkit.event.server.ServerListPingEvent;
 import org.bukkit.event.vehicle.VehicleCreateEvent;
 import org.bukkit.event.weather.LightningStrikeEvent;
+import org.bukkit.event.world.EntitiesLoadEvent;
+import org.bukkit.event.world.EntitiesUnloadEvent;
 import org.bukkit.event.world.LootGenerateEvent;
 import org.bukkit.inventory.EquipmentSlot;
 import org.bukkit.inventory.InventoryView;
@@ -1643,4 +1647,15 @@ public class CraftEventFactory {
         Bukkit.getPluginManager().callEvent(event);
         return event;
     }
+
+    public static void callEntitiesLoadEvent(World world, ChunkCoordIntPair coords, List<Entity> entities) {
+        List<org.bukkit.entity.Entity> bukkitEntities = Collections.unmodifiableList(entities.stream().map(Entity::getBukkitEntity).collect(Collectors.toList()));
+        EntitiesLoadEvent event = new EntitiesLoadEvent(new CraftChunk((WorldServer) world, coords.x, coords.z), bukkitEntities);
+        Bukkit.getPluginManager().callEvent(event);
+    }
+    public static void callEntitiesUnloadEvent(World world, ChunkCoordIntPair coords, List<Entity> entities) {
+        List<org.bukkit.entity.Entity> bukkitEntities = Collections.unmodifiableList(entities.stream().map(Entity::getBukkitEntity).collect(Collectors.toList()));
+        EntitiesUnloadEvent event = new EntitiesUnloadEvent(new CraftChunk((WorldServer) world, coords.x, coords.z), bukkitEntities);
+        Bukkit.getPluginManager().callEvent(event);
+    }
 }