[1.3] Fix query partition memory retention and failure diagnostics

.github/workflows/sonar-codecov.yml

@@ -43,6 +43,11 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+      - name: Set up JDK 21
+        uses: actions/setup-java@v4
+        with:
+          distribution: corretto
+          java-version: 21
       - name: Cache Maven packages
         uses: actions/cache@v4
         with:

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

@@ -34,6 +34,8 @@
  */
 class AutoResizingBuffer {
 
+  static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;
+
   private byte[] array;
   private int bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
   private final int initialCapacity;
@@ -59,16 +61,29 @@ public void resizeIfNecessary(int size) {
   }
 
   void ensureCapacity(int size) {
+    ensureCapacity(size, MAX_ARRAY_SIZE);
+  }
+
+  void ensureCapacity(int size, int maxCapacity) {
+    int limit = Math.min(maxCapacity, MAX_ARRAY_SIZE);
+    if (size < 0 || size > limit) {
+      throw new IllegalArgumentException("Invalid buffer size: " + size);
+    }
     final int currentCapacity = this.array.length;
     if (currentCapacity < size) {
-      // Increase by a factor of 1.5x
-      int growCapacity = currentCapacity + (currentCapacity >> 1);
-      int newCapacity = Math.max(growCapacity, size);
-      this.array = Arrays.copyOf(array, newCapacity);
+      this.array =
+          Arrays.copyOf(array, Math.min(calculateNewCapacity(currentCapacity, size), limit));
       bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
     }
   }
 
+  static int calculateNewCapacity(int currentCapacity, int requiredCapacity) {
+    // Keep geometric growth even above 1.33 GiB. An int overflow here would otherwise copy the
+    // entire buffer for every small write during Thrift serialization.
+    long growCapacity = currentCapacity + (long) (currentCapacity >> 1);
+    return (int) Math.min(MAX_ARRAY_SIZE, Math.max(growCapacity, requiredCapacity));
+  }
+
   public byte[] array() {
     return this.array;
   }

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

@@ -43,14 +43,41 @@ public int read(byte[] buf, int off, int len) {
   }
 
   @Override
-  public void write(byte[] toWrite, int off, int len) {
+  public void write(byte[] toWrite, int off, int len) throws TTransportException {
+    write(toWrite, off, len, AutoResizingBuffer.MAX_ARRAY_SIZE);
+  }
+
+  void write(byte[] toWrite, int off, int len, int maxCapacity) throws TTransportException {
+    int limit = Math.min(maxCapacity, AutoResizingBuffer.MAX_ARRAY_SIZE);
+    long requiredCapacity = (long) pos + len;
+    if (requiredCapacity > limit) {
+      throw new TTransportException(
+          TTransportException.CORRUPTED_DATA,
+          "Write buffer size ("
+              + requiredCapacity
+              + ") larger than maximum buffer size ("
+              + limit
+              + ")!");
+    }
     // A small prefix during serialization does not indicate a small frame. Only the complete
     // frame size, supplied after flushing, may trigger shrinking.
-    buf.ensureCapacity(pos + len);
+    buf.ensureCapacity((int) requiredCapacity, limit);
     System.arraycopy(toWrite, off, buf.array(), pos, len);
     pos += len;
   }
 
+  // Framed callers keep capacity within their limit. A fitting write needs no further limit
+  // checks; compare remaining space rather than adding pos and len, which could overflow.
+  boolean tryWriteWithinCapacity(byte[] toWrite, int off, int len) {
+    byte[] array = buf.array();
+    if (len > array.length - pos) {
+      return false;
+    }
+    System.arraycopy(toWrite, off, array, pos, len);
+    pos += len;
+    return true;
+  }
+
   public int getPos() {
     return pos;
   }

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

@@ -40,6 +40,13 @@ protected TCompressedElasticFramedTransport(
     readCompressBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize);
   }
 
+  @Override
+  protected int getMaxWriteBufferSize() {
+    // Preserve the wire-size limit: a compressible payload may exceed thriftMaxFrameSize before
+    // compression. Incremental writes still must fit in a Java array.
+    return AutoResizingBuffer.MAX_ARRAY_SIZE;
+  }
+
   @Override
   protected void readFrame() throws TTransportException {
     underlying.readAll(i32buf, 0, 4);
@@ -49,6 +56,7 @@ protected void readFrame() throws TTransportException {
     RpcStat.readCompressedBytes.addAndGet(size);
     try {
       int uncompressedLength = uncompressedLength(readBuffer.getBuffer(), 0, size);
+      checkBufferSize(uncompressedLength);
       RpcStat.readBytes.addAndGet(uncompressedLength);
       readCompressBuffer.resizeIfNecessary(uncompressedLength);
       uncompress(readBuffer.getBuffer(), 0, size, readCompressBuffer.getBuffer(), 0);
@@ -66,6 +74,7 @@ public void flush() throws TTransportException {
     RpcStat.writeBytes.addAndGet(length);
     try {
       int maxCompressedLength = maxCompressedLength(length);
+      checkBufferSize(maxCompressedLength);
       writeCompressBuffer.resizeIfNecessary(maxCompressedLength);
       int compressedLength =
           compress(writeBuffer.getBuffer(), 0, length, writeCompressBuffer.getBuffer(), 0);
@@ -83,6 +92,16 @@ public void flush() throws TTransportException {
     underlying.flush();
   }
 
+  private void checkBufferSize(int size) throws TTransportException {
+    if (size < 0 || size > AutoResizingBuffer.MAX_ARRAY_SIZE) {
+      // Compression bounds can overflow int even when the input fits in a Java array. Reject
+      // before resizing and invalidate the transport, just as for an oversized wire frame.
+      close();
+      throw new TTransportException(
+          TTransportException.CORRUPTED_DATA, "Invalid compression buffer size: " + size);
+    }
+  }
+
   protected abstract int uncompressedLength(byte[] but, int off, int len) throws IOException;
 
   protected abstract int maxCompressedLength(int len);

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

@@ -40,9 +40,8 @@ public class TElasticFramedTransport extends TTransport {
   public static class Factory extends TTransportFactory {
 
     /**
-     * It is used to prevent the size of the parsing package from being too large and allocating the
-     * buffer will cause oom. Therefore, the maximum length of the requested memory is limited when
-     * reading.
+     * Maximum size of a frame on the wire. Uncompressed writes are checked before buffering;
+     * compressed writes are checked after compression.
      */
     protected final int thriftMaxFrameSize;
 
@@ -88,7 +87,11 @@ public TElasticFramedTransport(
     this.thriftMaxFrameSize = thriftMaxFrameSize;
     this.copyBinary = copyBinary;
     readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize);
-    writeBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize);
+    writeBuffer =
+        new AutoScalingBufferWriteTransport(
+            Math.min(
+                thriftDefaultBufferSize,
+                Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE)));
   }
 
   protected final int thriftDefaultBufferSize;
@@ -288,8 +291,35 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {
   }
 
   @Override
-  public void write(byte[] buf, int off, int len) {
-    writeBuffer.write(buf, off, len);
+  public void write(byte[] buf, int off, int len) throws TTransportException {
+    if (writeBuffer.tryWriteWithinCapacity(buf, off, len)) {
+      return;
+    }
+    long size = (long) writeBuffer.getPos() + len;
+    int maxSize = getMaxWriteBufferSize();
+    if (size > maxSize) {
+      // Thrift serializes the whole request before flush. Reject it before growing the buffer so
+      // an oversized request cannot keep allocating and retain the caller's query memory.
+      close();
+      String limitDescription =
+          maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE
+              ? "maximum Java array size"
+              : "configured Thrift max frame size";
+      throw new TTransportException(
+          TTransportException.CORRUPTED_DATA,
+          "Serialized frame size is at least "
+              + size
+              + " bytes, exceeding the "
+              + limitDescription
+              + " of "
+              + maxSize
+              + " bytes.");
+    }
+    writeBuffer.write(buf, off, len, maxSize);
+  }
+
+  protected int getMaxWriteBufferSize() {
+    return Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE);
   }
 
   public TTransport getSocket() {

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

@@ -20,6 +20,7 @@
 package org.apache.iotdb.rpc;
 
 import org.apache.thrift.transport.TMemoryBuffer;
+import org.apache.thrift.transport.TTransportException;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -28,6 +29,47 @@
 
 public class AutoResizingBufferTest {
 
+  @Test
+  public void testGrowthNearIntegerLimitKeepsHeadroom() {
+    Assert.assertEquals(
+        AutoResizingBuffer.MAX_ARRAY_SIZE,
+        AutoResizingBuffer.calculateNewCapacity(1_500_000_000, 1_500_000_008));
+    Assert.assertEquals(
+        AutoResizingBuffer.MAX_ARRAY_SIZE,
+        AutoResizingBuffer.calculateNewCapacity(
+            AutoResizingBuffer.MAX_ARRAY_SIZE - 8, AutoResizingBuffer.MAX_ARRAY_SIZE));
+  }
+
+  @Test
+  public void testInvalidBufferSizeDoesNotResize() {
+    AutoResizingBuffer buffer = new AutoResizingBuffer(8);
+    for (int size : new int[] {-1, Integer.MAX_VALUE}) {
+      try {
+        buffer.ensureCapacity(size);
+        Assert.fail("The requested capacity must fit in a Java array");
+      } catch (IllegalArgumentException expected) {
+        Assert.assertEquals(8, buffer.array().length);
+      }
+    }
+  }
+
+  @Test
+  public void testWritePositionOverflowIsRejectedBeforeCopying() throws Exception {
+    AutoScalingBufferWriteTransport transport = new AutoScalingBufferWriteTransport(8);
+    Field position = AutoScalingBufferWriteTransport.class.getDeclaredField("pos");
+    position.setAccessible(true);
+    position.setInt(transport, Integer.MAX_VALUE - 4);
+    byte[] buffer = transport.getBuffer();
+    try {
+      transport.write(new byte[8]);
+      Assert.fail("The write position must not overflow");
+    } catch (TTransportException e) {
+      Assert.assertEquals(TTransportException.CORRUPTED_DATA, e.getType());
+      Assert.assertSame(buffer, transport.getBuffer());
+      Assert.assertEquals(Integer.MAX_VALUE - 4, transport.getPos());
+    }
+  }
+
   @Test
   public void testGrowthPreservesDataAndHeadroom() {
     AutoResizingBuffer buffer = new AutoResizingBuffer(100);