recycle BedrockPacketWrapper & add flag inheritance

CloudburstMC · Jul 22, 2024 · ab383d4 · ab383d4
1 parent 3b8a7c9
commit ab383d4
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 28 deletions.
diff --git a/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/BedrockPeer.java b/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/BedrockPeer.java
@@ -110,11 +110,11 @@ private void free() {
     }
 
     public void sendPacket(int senderClientId, int targetClientId, BedrockPacket packet) {
-        this.packetQueue.add(new BedrockPacketWrapper(0, senderClientId, targetClientId, packet, null));
+        this.packetQueue.add(BedrockPacketWrapper.create(0, senderClientId, targetClientId, packet, null));
     }
 
     public void sendPacketImmediately(int senderClientId, int targetClientId, BedrockPacket packet) {
-        this.channel.writeAndFlush(new BedrockPacketWrapper(0, senderClientId, targetClientId, packet, null));
+        this.channel.writeAndFlush(BedrockPacketWrapper.create(0, senderClientId, targetClientId, packet, null));
     }
 
     public void enableEncryption(@NonNull SecretKey secretKey) {

diff --git a/...connection/src/main/java/org/cloudburstmc/protocol/bedrock/netty/BedrockBatchWrapper.java b/...connection/src/main/java/org/cloudburstmc/protocol/bedrock/netty/BedrockBatchWrapper.java
@@ -11,7 +11,7 @@
 import lombok.EqualsAndHashCode;
 import org.cloudburstmc.protocol.bedrock.data.CompressionAlgorithm;
 import org.cloudburstmc.protocol.bedrock.packet.BedrockPacket;
-import org.cloudburstmc.protocol.bedrock.util.BatchFlag;
+import org.cloudburstmc.protocol.bedrock.util.PacketFlag;
 
 import java.util.List;
 import java.util.Set;
@@ -29,7 +29,7 @@ public class BedrockBatchWrapper extends AbstractReferenceCounted {
     private List<BedrockPacketWrapper> packets = new ObjectArrayList<>();
 
     private boolean modified;
-    private Set<BatchFlag> flags = new ObjectOpenHashSet<>();
+    private Set<PacketFlag> flags = new ObjectOpenHashSet<>();
 
     private BedrockBatchWrapper(ObjectPool.Handle<BedrockBatchWrapper> handle) {
         this.handle = handle;
@@ -41,20 +41,20 @@ public static BedrockBatchWrapper newInstance() {
 
     public static BedrockBatchWrapper newInstance(ByteBuf compressed, ByteBuf uncompressed) {
         BedrockBatchWrapper batch = RECYCLER.get();
-        batch.compressed = compressed;
-        batch.uncompressed = uncompressed;
-        batch.setRefCnt(1);
-
         if (!batch.packets.isEmpty() || !batch.flags.isEmpty()) {
             throw new IllegalStateException("Batch was not deallocated");
         }
+
+        batch.compressed = compressed;
+        batch.uncompressed = uncompressed;
+        batch.setRefCnt(1);
         return batch;
     }
 
     public static BedrockBatchWrapper create(int subClientId, BedrockPacket... packets) {
         BedrockBatchWrapper batch = BedrockBatchWrapper.newInstance();
         for (BedrockPacket packet : packets) {
-            batch.getPackets().add(new BedrockPacketWrapper(0, subClientId, 0, packet, null));
+            batch.getPackets().add(BedrockPacketWrapper.create(0, subClientId, 0, packet, null));
         }
         return batch;
     }
@@ -76,6 +76,14 @@ protected void deallocate() {
     public void addPacket(BedrockPacketWrapper wrapper) {
         this.packets.add(wrapper);
         this.modify();
+
+        if (!wrapper.getFlags().isEmpty()) {
+            for (PacketFlag flag : wrapper.getFlags()) {
+                if (flag.canInherit()) {
+                    this.flags.add(flag);
+                }
+            }
+        }
     }
 
     public void modify() {
@@ -109,15 +117,15 @@ public void setUncompressed(ByteBuf uncompressed) {
         this.uncompressed = uncompressed;
     }
 
-    public void setFlag(BatchFlag flag) {
+    public void setFlag(PacketFlag flag) {
         this.flags.add(flag);
     }
 
-    public boolean hasFlag(BatchFlag flag) {
+    public boolean hasFlag(PacketFlag flag) {
         return this.flags.contains(flag);
     }
 
-    public void unsetFlag(BatchFlag flag) {
+    public void unsetFlag(PacketFlag flag) {
         this.flags.remove(flag);
     }
 

diff --git a/...onnection/src/main/java/org/cloudburstmc/protocol/bedrock/netty/BedrockPacketWrapper.java b/...onnection/src/main/java/org/cloudburstmc/protocol/bedrock/netty/BedrockPacketWrapper.java
@@ -3,37 +3,83 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.util.AbstractReferenceCounted;
 import io.netty.util.ReferenceCountUtil;
-import io.netty.util.ReferenceCounted;
-import lombok.AllArgsConstructor;
+import io.netty.util.internal.ObjectPool;
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
-import lombok.NoArgsConstructor;
 import org.cloudburstmc.protocol.bedrock.packet.BedrockPacket;
+import org.cloudburstmc.protocol.bedrock.util.PacketFlag;
+
+import java.util.Set;
 
 @Data
 @EqualsAndHashCode(callSuper = false)
-@AllArgsConstructor
-@NoArgsConstructor
 public class BedrockPacketWrapper extends AbstractReferenceCounted {
+    private static final ObjectPool<BedrockPacketWrapper> RECYCLER = ObjectPool.newPool(BedrockPacketWrapper::new);
+    private final ObjectPool.Handle<BedrockPacketWrapper> handle;
+
     private int packetId;
     private int senderSubClientId;
     private int targetSubClientId;
     private int headerLength;
     private BedrockPacket packet;
     private ByteBuf packetBuffer;
+    private Set<PacketFlag> flags = new ObjectOpenHashSet<>();
+
+    public static BedrockPacketWrapper create(int packetId, int senderSubClientId, int targetSubClientId, BedrockPacket packet, ByteBuf packetBuffer) {
+        BedrockPacketWrapper wrapper = RECYCLER.get();
+        if (wrapper.packet != null || wrapper.packetBuffer != null) {
+            throw new IllegalStateException("BedrockPacketWrapper was not deallocated");
+        }
+
+        wrapper.packetId = packetId;
+        wrapper.senderSubClientId = senderSubClientId;
+        wrapper.targetSubClientId = targetSubClientId;
+        wrapper.packet = packet;
+        wrapper.packetBuffer = packetBuffer;
+
+        wrapper.setRefCnt(1);
+        return wrapper;
+    }
+
+    public static BedrockPacketWrapper create() {
+        BedrockPacketWrapper wrapper = RECYCLER.get();
+        if (wrapper.packet != null || wrapper.packetBuffer != null) {
+            throw new IllegalStateException("BedrockPacketWrapper was not deallocated");
+        }
+
+        wrapper.setRefCnt(1);
+        return wrapper;
+    }
+
+    private BedrockPacketWrapper(ObjectPool.Handle<BedrockPacketWrapper> handle) {
+        this.handle = handle;
+    }
+
+    public void setFlag(PacketFlag flag) {
+        this.flags.add(flag);
+    }
+
+    public boolean hasFlag(PacketFlag flag) {
+        return this.flags.contains(flag);
+    }
 
-    public BedrockPacketWrapper(int packetId, int senderSubClientId, int targetSubClientId, BedrockPacket packet, ByteBuf packetBuffer) {
-        this.packetId = packetId;
-        this.senderSubClientId = senderSubClientId;
-        this.targetSubClientId = targetSubClientId;
-        this.packet = packet;
-        this.packetBuffer = packetBuffer;
+    public void unsetFlag(PacketFlag flag) {
+        this.flags.remove(flag);
     }
 
     @Override
     protected void deallocate() {
         ReferenceCountUtil.safeRelease(this.packet);
         ReferenceCountUtil.safeRelease(this.packetBuffer);
+        this.packetId = 0;
+        this.senderSubClientId = 0;
+        this.targetSubClientId = 0;
+        this.headerLength = 0;
+        this.packet = null;
+        this.packetBuffer = null;
+        this.flags.clear();
+        this.handle.recycle(this);
     }
 
     @Override

diff --git a/...rc/main/java/org/cloudburstmc/protocol/bedrock/netty/codec/packet/BedrockPacketCodec.java b/...rc/main/java/org/cloudburstmc/protocol/bedrock/netty/codec/packet/BedrockPacketCodec.java
@@ -62,7 +62,7 @@ protected final void encode(ChannelHandlerContext ctx, BedrockPacketWrapper msg,
 
     @Override
     protected final void decode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
-        BedrockPacketWrapper wrapper = new BedrockPacketWrapper();
+        BedrockPacketWrapper wrapper = BedrockPacketWrapper.create();
         wrapper.setPacketBuffer(msg.retainedSlice());
         try {
             int index = msg.readerIndex();

diff --git a/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/util/BatchFlag.java b/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/util/BatchFlag.java
diff --git a/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/util/PacketFlag.java b/bedrock-connection/src/main/java/org/cloudburstmc/protocol/bedrock/util/PacketFlag.java
@@ -0,0 +1,12 @@
+package org.cloudburstmc.protocol.bedrock.util;
+
+public interface PacketFlag {
+
+    /**
+     * Whether this flag can be passed from packet wrapper to batch wrapper.
+     * @return true if the flag can be passed
+     */
+    default boolean canInherit() {
+        return true;
+    }
+}