Reset message budget for compressed frames

Author: HTHou

iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java

@@ -59,10 +59,14 @@ protected void closeAllocatedBuffers() {
 
   @Override
   protected void readFrame() throws TTransportException {
+    // Discard the previous frame's budget before reading the next frame header and payload.
+    resetMessageSizeAndConsumedBytes();
     underlying.readAll(i32buf, 0, 4);
     int size = TFramedTransport.decodeFrameSize(i32buf);
     validateFrame(size);
     readBuffer.fill(underlying, size);
+    // Bind subsequent protocol reads to the current compressed frame size.
+    resetMessageSizeAndConsumedBytes(size);
     RpcStat.readCompressedBytes.addAndGet(size);
     try {
       int uncompressedLength = uncompressedLength(readBuffer.getBuffer(), 0, size);

iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java

@@ -27,6 +27,7 @@
 
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -61,6 +62,44 @@ public void testReadFramesResetMessageSize() throws TTransportException {
     assertArrayEquals(secondFrame, actualSecondFrame);
   }
 
+  @Test
+  public void testReadCompressedFramesResetMessageSize() throws TTransportException {
+    byte[] firstFrame = {1, 2};
+    byte[] secondFrame = {3, 4, 5, 6, 7, 8};
+    TMemoryBuffer wire = new TMemoryBuffer(128);
+    TSnappyElasticFramedTransport output =
+        new TSnappyElasticFramedTransport(wire, 4, 128, false);
+    output.write(firstFrame);
+    output.flush();
+    output.write(secondFrame);
+    output.flush();
+
+    byte[] framedData = Arrays.copyOf(wire.getArray(), wire.length());
+    ByteBuffer frameSizes = ByteBuffer.wrap(framedData);
+    int firstCompressedSize = frameSizes.getInt();
+    frameSizes.position(Integer.BYTES + firstCompressedSize);
+    int secondCompressedSize = frameSizes.getInt();
+    int maxMessageSize = Integer.BYTES + Math.max(firstCompressedSize, secondCompressedSize);
+    TConfiguration configuration =
+        TConfiguration.custom()
+            .setMaxMessageSize(maxMessageSize)
+            .setMaxFrameSize(maxMessageSize)
+            .build();
+    TMemoryBuffer underlying = new TMemoryBuffer(configuration, maxMessageSize);
+    underlying.write(framedData);
+    TSnappyElasticFramedTransport input =
+        new TSnappyElasticFramedTransport(
+            underlying, 4, configuration.getMaxFrameSize(), false);
+
+    byte[] actualFirstFrame = new byte[firstFrame.length];
+    input.readAll(actualFirstFrame, 0, actualFirstFrame.length);
+    assertArrayEquals(firstFrame, actualFirstFrame);
+
+    byte[] actualSecondFrame = new byte[secondFrame.length];
+    input.readAll(actualSecondFrame, 0, actualSecondFrame.length);
+    assertArrayEquals(secondFrame, actualSecondFrame);
+  }
+
   @Test
   public void testSingularSize() {