fix async chunk sending buffer (#301)

* fix Async chunk sending buffer size

* cleanup

* synchronized write nonEmptyBlockCount

* increase buffer initial capacity

* fix block entity map data race

---------

Co-authored-by: Taiyou06 <kaandindar21@gmail.com>

leaf-server/minecraft-patches/features/0132-Async-chunk-sending.patch

@@ -5,10 +5,10 @@ Subject: [PATCH] Async chunk sending
 
 
 diff --git a/ca/spottedleaf/moonrise/patches/chunk_system/player/RegionizedPlayerChunkLoader.java b/ca/spottedleaf/moonrise/patches/chunk_system/player/RegionizedPlayerChunkLoader.java
-index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2fccbbc5e6 100644
+index a35e9fae8f8da0c42f0616c4f78dc396492673aa..c0f14d2bc30a186511bafddc2e80f29b686887c5 100644
 --- a/ca/spottedleaf/moonrise/patches/chunk_system/player/RegionizedPlayerChunkLoader.java
 +++ b/ca/spottedleaf/moonrise/patches/chunk_system/player/RegionizedPlayerChunkLoader.java
-@@ -411,19 +411,91 @@ public final class RegionizedPlayerChunkLoader {
+@@ -411,19 +411,98 @@ public final class RegionizedPlayerChunkLoader {
              this.delayedTicketOps.addLast(op);
          }
  
@@ -29,7 +29,6 @@ index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2f
 +                // Already in our sent list - silently return instead of throwing an exception
 +                return;
 +            }
-+
 +            // Get the chunk now, as we need it for both sync and async paths
 +            final LevelChunk chunk = ((ChunkSystemLevel) this.world).moonrise$getFullChunkIfLoaded(chunkX, chunkZ);
 +            if (chunk == null) {
@@ -37,13 +36,12 @@ index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2f
 +                this.sentChunks.remove(chunkKey);
 +                return;
 +            }
-+
+ 
 +            // Try to mark the chunk as received by this player
 +            try {
 +                // This part needs to remain on the main thread as it affects shared state
 +                ((ChunkSystemServerLevel) this.world).moonrise$getChunkTaskScheduler().chunkHolderManager
 +                    .getChunkHolder(chunkX, chunkZ).vanillaChunkHolder.moonrise$addReceivedChunk(this.player);
- 
 +                // Call onChunkWatch on the main thread as it might affect server state
                  PlatformHooks.get().onChunkWatch(this.world, chunk, this.player);
 -                PlayerChunkSender.sendChunk(this.player.connection, this.world, chunk);
@@ -58,6 +56,13 @@ index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2f
 +            // Check if async chunk sending is enabled
 +            if (org.dreeam.leaf.config.modules.async.AsyncChunkSend.enabled) {
 +                // Async implementation
++                var heightmaps = new net.minecraft.nbt.CompoundTag();
++                for (var entry : chunk.getHeightmaps()) {
++                    if (entry.getKey().sendToClient()) {
++                        heightmaps.put(entry.getKey().getSerializationKey(), new net.minecraft.nbt.LongArrayTag(entry.getValue().getRawData()));
++                    }
++                }
++                var blockEntities = chunk.blockEntities.values().toArray(new net.minecraft.world.level.block.entity.BlockEntity[0]);
 +                net.minecraft.Util.backgroundExecutor().execute(() -> {
 +                    try {
 +                        final net.minecraft.server.network.ServerGamePacketListenerImpl connection = this.player.connection;
@@ -66,7 +71,9 @@ index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2f
 +                        // Create the packet with anti-xray control flag
 +                        final net.minecraft.network.protocol.game.ClientboundLevelChunkWithLightPacket packet = new net.minecraft.network.protocol.game.ClientboundLevelChunkWithLightPacket(
 +                            chunk, serverLevel.getLightEngine(), null, null,
-+                            serverLevel.chunkPacketBlockController.shouldModify(this.player, chunk)
++                            serverLevel.chunkPacketBlockController.shouldModify(this.player, chunk),
++                            heightmaps,
++                            blockEntities
 +                        );
 +
 +                        // Let the main thread handle the anti-xray processing
@@ -106,3 +113,138 @@ index a35e9fae8f8da0c42f0616c4f78dc396492673aa..af49117695c0785033b984ff91550e2f
  
          private void sendUnloadChunk(final int chunkX, final int chunkZ) {
              if (!this.sentChunks.remove(CoordinateUtils.getChunkKey(chunkX, chunkZ))) {
+diff --git a/net/minecraft/network/protocol/game/ClientboundLevelChunkPacketData.java b/net/minecraft/network/protocol/game/ClientboundLevelChunkPacketData.java
+index 9e321ef1c3d5803519b243685f4ee598dc0cf640..c981801171307e571388e6e810840ae2525c5d10 100644
+--- a/net/minecraft/network/protocol/game/ClientboundLevelChunkPacketData.java
++++ b/net/minecraft/network/protocol/game/ClientboundLevelChunkPacketData.java
+@@ -26,6 +26,7 @@ public class ClientboundLevelChunkPacketData {
+     private static final int TWO_MEGABYTES = 2097152;
+     private final CompoundTag heightmaps;
+     private final byte[] buffer;
++    private final int bufferLength;
+     private final List<ClientboundLevelChunkPacketData.BlockEntityInfo> blockEntitiesData;
+     // Paper start - Handle oversized block entities in chunks
+     private final java.util.List<net.minecraft.network.protocol.Packet<?>> extraPackets = new java.util.ArrayList<>();
+@@ -52,6 +53,7 @@ public class ClientboundLevelChunkPacketData {
+         }
+ 
+         this.buffer = new byte[calculateChunkSize(levelChunk)];
++        this.bufferLength = this.buffer.length; // Leaf
+         // Paper start - Anti-Xray - Add chunk packet info
+         if (chunkPacketInfo != null) {
+             chunkPacketInfo.setBuffer(this.buffer);
+@@ -74,6 +76,51 @@ public class ClientboundLevelChunkPacketData {
+         }
+     }
+ 
++    // Leaf start - Async chunk sending
++    public ClientboundLevelChunkPacketData(
++        LevelChunk levelChunk,
++        io.papermc.paper.antixray.ChunkPacketInfo<net.minecraft.world.level.block.state.BlockState> chunkPacketInfo,
++        CompoundTag heightmaps,
++        BlockEntity[] blockEntities
++    ) {
++        this.heightmaps = heightmaps;
++        var buffer = new FriendlyByteBuf(Unpooled.buffer(calculateChunkSize(levelChunk) + 128));
++        var sections = levelChunk.getSections();
++        var sectionLength = sections.length;
++        for (int i = 0; i < sectionLength; i++) {
++            LevelChunkSection section = sections[i];
++            synchronized (section.getStates()) {
++                buffer.writeShort(section.nonEmptyBlockCount());
++                section.getStates().write(buffer, chunkPacketInfo, i);
++            }
++            section.getBiomes().write(buffer, null, i);
++        }
++        this.buffer = buffer.array();
++        this.bufferLength = buffer.writerIndex();
++
++        // Paper start - Anti-Xray - Add chunk packet info
++        if (chunkPacketInfo != null) {
++            chunkPacketInfo.setBuffer(this.buffer);
++            chunkPacketInfo.setLength(this.bufferLength);
++        }
++        this.blockEntitiesData = Lists.newArrayList();
++        int totalTileEntities = 0; // Paper - Handle oversized block entities in chunks
++
++        for (BlockEntity blockEntity : blockEntities) {
++            // Paper start - Handle oversized block entities in chunks
++            if (++totalTileEntities > BLOCK_ENTITY_LIMIT) {
++                net.minecraft.network.protocol.Packet<ClientGamePacketListener> packet = blockEntity.getUpdatePacket();
++                if (packet != null) {
++                    this.extraPackets.add(packet);
++                    continue;
++                }
++            }
++            // Paper end - Handle oversized block entities in chunks
++            this.blockEntitiesData.add(ClientboundLevelChunkPacketData.BlockEntityInfo.create(blockEntity));
++        }
++    }
++    // Leaf end - Async chunk sending
++
+     public ClientboundLevelChunkPacketData(RegistryFriendlyByteBuf buffer, int x, int z) {
+         this.heightmaps = buffer.readNbt();
+         if (this.heightmaps == null) {
+@@ -84,6 +131,7 @@ public class ClientboundLevelChunkPacketData {
+                 throw new RuntimeException("Chunk Packet trying to allocate too much memory on read.");
+             } else {
+                 this.buffer = new byte[varInt];
++                this.bufferLength = this.buffer.length; // Leaf
+                 buffer.readBytes(this.buffer);
+                 this.blockEntitiesData = ClientboundLevelChunkPacketData.BlockEntityInfo.LIST_STREAM_CODEC.decode(buffer);
+             }
+@@ -92,8 +140,8 @@ public class ClientboundLevelChunkPacketData {
+ 
+     public void write(RegistryFriendlyByteBuf buffer) {
+         buffer.writeNbt(this.heightmaps);
+-        buffer.writeVarInt(this.buffer.length);
+-        buffer.writeBytes(this.buffer);
++        buffer.writeVarInt(this.bufferLength); // Leaf
++        buffer.writeBytes(this.buffer, 0, this.bufferLength); // Leaf
+         ClientboundLevelChunkPacketData.BlockEntityInfo.LIST_STREAM_CODEC.encode(buffer, this.blockEntitiesData);
+     }
+ 
+diff --git a/net/minecraft/network/protocol/game/ClientboundLevelChunkWithLightPacket.java b/net/minecraft/network/protocol/game/ClientboundLevelChunkWithLightPacket.java
+index 8578d1f78ddd1bb75f3230f04bfaa35af9f5f822..4eeb4967120b1c2cf13d2b4a8c07175fb4d98012 100644
+--- a/net/minecraft/network/protocol/game/ClientboundLevelChunkWithLightPacket.java
++++ b/net/minecraft/network/protocol/game/ClientboundLevelChunkWithLightPacket.java
+@@ -45,6 +45,26 @@ public class ClientboundLevelChunkWithLightPacket implements Packet<ClientGamePa
+         chunk.getLevel().chunkPacketBlockController.modifyBlocks(this, chunkPacketInfo); // Paper - Anti-Xray - Modify blocks
+     }
+ 
++    // Leaf start - Async chunk sending
++    public ClientboundLevelChunkWithLightPacket(LevelChunk chunk,
++                                                LevelLightEngine lightEngine,
++                                                @Nullable BitSet skyLight,
++                                                @Nullable BitSet blockLight,
++                                                boolean modifyBlocks,
++                                                net.minecraft.nbt.CompoundTag heightmaps,
++                                                net.minecraft.world.level.block.entity.BlockEntity[] blockEntities
++    ) {
++        // Paper end - Anti-Xray
++        ChunkPos pos = chunk.getPos();
++        this.x = pos.x;
++        this.z = pos.z;
++        io.papermc.paper.antixray.ChunkPacketInfo<net.minecraft.world.level.block.state.BlockState> chunkPacketInfo = modifyBlocks ? chunk.getLevel().chunkPacketBlockController.getChunkPacketInfo(this, chunk) : null; // Paper - Ant-Xray
++        this.chunkData = new ClientboundLevelChunkPacketData(chunk, chunkPacketInfo, heightmaps, blockEntities); // Paper - Anti-Xray
++        this.lightData = new ClientboundLightUpdatePacketData(pos, lightEngine, skyLight, blockLight);
++        chunk.getLevel().chunkPacketBlockController.modifyBlocks(this, chunkPacketInfo); // Paper - Anti-Xray - Modify blocks
++    }
++    // Leaf end - Async chunk sending
++
+     private ClientboundLevelChunkWithLightPacket(RegistryFriendlyByteBuf buffer) {
+         this.x = buffer.readInt();
+         this.z = buffer.readInt();
+diff --git a/net/minecraft/world/level/chunk/LevelChunkSection.java b/net/minecraft/world/level/chunk/LevelChunkSection.java
+index b8ac6a9ba7b56ccd034757f7d135d272b8e69e90..f85c1c3ab8cc0d0dafa6df8536318e3443265d2b 100644
+--- a/net/minecraft/world/level/chunk/LevelChunkSection.java
++++ b/net/minecraft/world/level/chunk/LevelChunkSection.java
+@@ -49,6 +49,8 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
+     }
+     // Paper end - block counting
+ 
++    public final int nonEmptyBlockCount() { return this.nonEmptyBlockCount; } // Leaf
++
+     private LevelChunkSection(LevelChunkSection section) {
+         this.nonEmptyBlockCount = section.nonEmptyBlockCount;
+         this.tickingBlockCount = section.tickingBlockCount;

leaf-server/minecraft-patches/features/0157-Optimise-chunkUnloads.patch

@@ -159,7 +159,7 @@ index 4ca68a903e67606fc4ef0bfa9862a73797121c8b..bed3a64388bb43e47c2ba4e67f7dde5b
  
      public static final class SaveState {
 diff --git a/net/minecraft/world/level/chunk/LevelChunkSection.java b/net/minecraft/world/level/chunk/LevelChunkSection.java
-index b8ac6a9ba7b56ccd034757f7d135d272b8e69e90..dc158e981199b507531af810ff9ced3ca717e39e 100644
+index f85c1c3ab8cc0d0dafa6df8536318e3443265d2b..449f180c49f191786c5f253baf04ddd9b72835f7 100644
 --- a/net/minecraft/world/level/chunk/LevelChunkSection.java
 +++ b/net/minecraft/world/level/chunk/LevelChunkSection.java
 @@ -24,6 +24,7 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
@@ -170,7 +170,7 @@ index b8ac6a9ba7b56ccd034757f7d135d272b8e69e90..dc158e981199b507531af810ff9ced3c
  
      // Paper start - block counting
      private static final it.unimi.dsi.fastutil.shorts.ShortArrayList FULL_LIST = new it.unimi.dsi.fastutil.shorts.ShortArrayList(16*16*16);
-@@ -135,6 +136,7 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
+@@ -137,6 +138,7 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
      // Paper end - block counting
  
      public BlockState setBlockState(int x, int y, int z, BlockState state, boolean useLocks) {
@@ -178,7 +178,7 @@ index b8ac6a9ba7b56ccd034757f7d135d272b8e69e90..dc158e981199b507531af810ff9ced3c
          BlockState blockState;
          if (useLocks) {
              blockState = this.states.getAndSet(x, y, z, state);
-@@ -328,7 +330,32 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
+@@ -330,7 +332,32 @@ public class LevelChunkSection implements ca.spottedleaf.moonrise.patches.block_
          this.biomes = palettedContainer;
      }
  

leaf-server/paper-patches/features/0034-Async-chunk-sending.patch

@@ -0,0 +1,124 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: hayanesuru <mc@jvavav.com>
+Date: Sun, 27 Apr 2025 21:19:30 +0800
+Subject: [PATCH] Async chunk sending
+
+
+diff --git a/src/main/java/io/papermc/paper/antixray/BitStorageReader.java b/src/main/java/io/papermc/paper/antixray/BitStorageReader.java
+index c27703775c845a8b0bd1b0cb8f05eb736d8a813c..7e72c910232b85b411ecba1fa0acdf7f81c85418 100644
+--- a/src/main/java/io/papermc/paper/antixray/BitStorageReader.java
++++ b/src/main/java/io/papermc/paper/antixray/BitStorageReader.java
+@@ -8,11 +8,18 @@ public final class BitStorageReader {
+     private int longInBufferIndex;
+     private int bitInLongIndex;
+     private long current;
++    private int length; // Leaf
+ 
+     public void setBuffer(byte[] buffer) {
+         this.buffer = buffer;
+     }
+ 
++    // Leaf start
++    public void setLength(int length) {
++        this.length = length;
++    }
++    // Leaf end
++
+     public void setBits(int bits) {
+         this.bits = bits;
+         mask = (1 << bits) - 1;
+@@ -25,7 +32,7 @@ public final class BitStorageReader {
+     }
+ 
+     private void init() {
+-        if (buffer.length > longInBufferIndex + 7) {
++        if (length > longInBufferIndex + 7) { // Leaf
+             current = ((((long) buffer[longInBufferIndex]) << 56)
+                 | (((long) buffer[longInBufferIndex + 1] & 0xff) << 48)
+                 | (((long) buffer[longInBufferIndex + 2] & 0xff) << 40)
+diff --git a/src/main/java/io/papermc/paper/antixray/BitStorageWriter.java b/src/main/java/io/papermc/paper/antixray/BitStorageWriter.java
+index 83412e0ddaade11eb7ac7b41bb8ae5b085802775..d521bd2275152575f5fe5038a817471026abccda 100644
+--- a/src/main/java/io/papermc/paper/antixray/BitStorageWriter.java
++++ b/src/main/java/io/papermc/paper/antixray/BitStorageWriter.java
+@@ -9,11 +9,18 @@ public final class BitStorageWriter {
+     private int bitInLongIndex;
+     private long current;
+     private boolean dirty;
++    private int length; // Leaf
+ 
+     public void setBuffer(byte[] buffer) {
+         this.buffer = buffer;
+     }
+ 
++    // Leaf start
++    public void setLength(int length) {
++        this.length = length;
++    }
++    // Leaf end
++
+     public void setBits(int bits) {
+         this.bits = bits;
+         mask = (1L << bits) - 1;
+@@ -26,7 +33,7 @@ public final class BitStorageWriter {
+     }
+ 
+     private void init() {
+-        if (buffer.length > longInBufferIndex + 7) {
++        if (length > longInBufferIndex + 7) { // Leaf
+             current = ((((long) buffer[longInBufferIndex]) << 56)
+                 | (((long) buffer[longInBufferIndex + 1] & 0xff) << 48)
+                 | (((long) buffer[longInBufferIndex + 2] & 0xff) << 40)
+@@ -41,7 +48,7 @@ public final class BitStorageWriter {
+     }
+ 
+     public void flush() {
+-        if (dirty && buffer.length > longInBufferIndex + 7) {
++        if (dirty && length > longInBufferIndex + 7) { // Leaf
+             buffer[longInBufferIndex] = (byte) (current >> 56 & 0xff);
+             buffer[longInBufferIndex + 1] = (byte) (current >> 48 & 0xff);
+             buffer[longInBufferIndex + 2] = (byte) (current >> 40 & 0xff);
+diff --git a/src/main/java/io/papermc/paper/antixray/ChunkPacketBlockControllerAntiXray.java b/src/main/java/io/papermc/paper/antixray/ChunkPacketBlockControllerAntiXray.java
+index ee2d3a54d760f9c26542eab03c51651a30e279a0..bbf338e5e98a7b7d1c7cd986333df03f4aad4bb1 100644
+--- a/src/main/java/io/papermc/paper/antixray/ChunkPacketBlockControllerAntiXray.java
++++ b/src/main/java/io/papermc/paper/antixray/ChunkPacketBlockControllerAntiXray.java
+@@ -227,6 +227,8 @@ public final class ChunkPacketBlockControllerAntiXray extends ChunkPacketBlockCo
+         boolean[] obfuscateTemp = null;
+         bitStorageReader.setBuffer(chunkPacketInfoAntiXray.getBuffer());
+         bitStorageWriter.setBuffer(chunkPacketInfoAntiXray.getBuffer());
++        bitStorageReader.setLength(chunkPacketInfoAntiXray.getLength());
++        bitStorageWriter.setLength(chunkPacketInfoAntiXray.getLength());
+         int numberOfBlocks = presetBlockStateBits.length;
+         // Keep the lambda expressions as simple as possible. They are used very frequently.
+         LayeredIntSupplier random = numberOfBlocks == 1 ? (() -> 0) : engineMode == EngineMode.OBFUSCATE_LAYER ? new LayeredIntSupplier() {
+diff --git a/src/main/java/io/papermc/paper/antixray/ChunkPacketInfo.java b/src/main/java/io/papermc/paper/antixray/ChunkPacketInfo.java
+index a33a4d45d478ededff27244fcb910d3f369f2151..26f0e4b43c2adc27e0ca2bafb7a1bab3fcf8f6cc 100644
+--- a/src/main/java/io/papermc/paper/antixray/ChunkPacketInfo.java
++++ b/src/main/java/io/papermc/paper/antixray/ChunkPacketInfo.java
+@@ -13,6 +13,7 @@ public class ChunkPacketInfo<T> {
+     private final int[] indexes;
+     private final Object[][] presetValues;
+     private byte[] buffer;
++    private int length; // Leaf
+ 
+     public ChunkPacketInfo(ClientboundLevelChunkWithLightPacket chunkPacket, LevelChunk chunk) {
+         this.chunkPacket = chunkPacket;
+@@ -36,8 +37,19 @@ public class ChunkPacketInfo<T> {
+         return buffer;
+     }
+ 
++    // Leaf start
++    public int getLength() {
++        return length;
++    }
++    // Leaf end
++
+     public void setBuffer(byte[] buffer) {
+         this.buffer = buffer;
++        this.length = buffer.length; // Leaf
++    }
++
++    public void setLength(int length) {
++        this.length = length; // Leaf
+     }
+ 
+     public int getBits(int chunkSectionIndex) {