Align framed transport message budgets

Author: HTHou

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

@@ -91,6 +91,7 @@ public TElasticFramedTransport(
     this.thriftDefaultBufferSize = thriftDefaultBufferSize;
     this.thriftMaxFrameSize = thriftMaxFrameSize;
     this.copyBinary = copyBinary;
+    alignUnderlyingConfiguration();
     try {
       readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize);
       writeBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize);
@@ -100,6 +101,25 @@ public TElasticFramedTransport(
     }
   }
 
+  private void alignUnderlyingConfiguration() throws TTransportException {
+    TConfiguration configuration = underlying.getConfiguration();
+    if (configuration == null) {
+      return;
+    }
+    long maxMessageSize = (long) thriftMaxFrameSize + Integer.BYTES;
+    if (maxMessageSize > Integer.MAX_VALUE) {
+      throw new TTransportException(
+          TTransportException.MESSAGE_SIZE_LIMIT,
+          "Frame size exceeds the maximum supported message size: " + thriftMaxFrameSize);
+    }
+    if (configuration.getMaxFrameSize() < thriftMaxFrameSize) {
+      configuration.setMaxFrameSize(thriftMaxFrameSize);
+    }
+    if (configuration.getMaxMessageSize() < maxMessageSize) {
+      configuration.setMaxMessageSize((int) maxMessageSize);
+    }
+  }
+
   protected final int thriftDefaultBufferSize;
   protected final int thriftMaxFrameSize;
 

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

@@ -47,8 +47,8 @@ public void testReadFramesResetMessageSize() throws TTransportException {
             .put(secondFrame)
             .flip();
     TConfiguration configuration =
-        TConfiguration.custom().setMaxMessageSize(10).setMaxFrameSize(10).build();
-    TMemoryBuffer underlying = new TMemoryBuffer(configuration, 10);
+        TConfiguration.custom().setMaxMessageSize(14).setMaxFrameSize(10).build();
+    TMemoryBuffer underlying = new BudgetConsumingMemoryBuffer(configuration, 14);
     underlying.write(framedData.array());
     TElasticFramedTransport transport =
         new TElasticFramedTransport(underlying, 4, configuration.getMaxFrameSize(), false);
@@ -78,16 +78,17 @@ public void testReadCompressedFramesResetMessageSize() throws TTransportExceptio
     int firstCompressedSize = frameSizes.getInt();
     frameSizes.position(Integer.BYTES + firstCompressedSize);
     int secondCompressedSize = frameSizes.getInt();
-    int maxMessageSize = Integer.BYTES + Math.max(firstCompressedSize, secondCompressedSize);
+    int maxFrameSize = Math.max(firstCompressedSize, secondCompressedSize);
     TConfiguration configuration =
         TConfiguration.custom()
-            .setMaxMessageSize(maxMessageSize)
-            .setMaxFrameSize(maxMessageSize)
+            .setMaxMessageSize(maxFrameSize)
+            .setMaxFrameSize(maxFrameSize)
             .build();
-    TMemoryBuffer underlying = new TMemoryBuffer(configuration, maxMessageSize);
+    TMemoryBuffer underlying = new BudgetConsumingMemoryBuffer(configuration, maxFrameSize);
     underlying.write(framedData);
     TSnappyElasticFramedTransport input =
         new TSnappyElasticFramedTransport(underlying, 4, configuration.getMaxFrameSize(), false);
+    assertEquals(Integer.BYTES + maxFrameSize, configuration.getMaxMessageSize());
 
     byte[] actualFirstFrame = new byte[firstFrame.length];
     input.readAll(actualFirstFrame, 0, actualFirstFrame.length);
@@ -98,6 +99,39 @@ public void testReadCompressedFramesResetMessageSize() throws TTransportExceptio
     assertArrayEquals(secondFrame, actualSecondFrame);
   }
 
+  @Test
+  public void testFrameSizeAlignsUnderlyingMessageSize() throws TTransportException {
+    byte[] frame = {1, 2, 3, 4, 5, 6, 7, 8};
+    ByteBuffer framedData =
+        ByteBuffer.allocate(Integer.BYTES + frame.length).putInt(frame.length).put(frame).flip();
+    TConfiguration configuration =
+        TConfiguration.custom().setMaxMessageSize(6).setMaxFrameSize(frame.length).build();
+    TMemoryBuffer underlying = new BudgetConsumingMemoryBuffer(configuration, 6);
+    underlying.write(framedData.array());
+    TElasticFramedTransport transport =
+        new TElasticFramedTransport(underlying, 4, frame.length, false);
+
+    assertEquals(Integer.BYTES + frame.length, configuration.getMaxMessageSize());
+    byte[] actualFrame = new byte[frame.length];
+    transport.readAll(actualFrame, 0, actualFrame.length);
+    assertArrayEquals(frame, actualFrame);
+  }
+
+  private static class BudgetConsumingMemoryBuffer extends TMemoryBuffer {
+
+    private BudgetConsumingMemoryBuffer(TConfiguration configuration, int size)
+        throws TTransportException {
+      super(configuration, size);
+    }
+
+    @Override
+    public int read(byte[] buffer, int offset, int length) throws TTransportException {
+      int bytesRead = super.read(buffer, offset, length);
+      countConsumedMessageBytes(bytesRead);
+      return bytesRead;
+    }
+  }
+
   @Test
   public void testSingularSize() {