Fix query partition memory retention and failure diagnostics

iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java

@@ -127,4 +127,25 @@ private RpcMessages() {}
 
   public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 =
       "Data type %s is not supported.";
+  public static final String EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA =
+      "Invalid buffer size: %d";
+
+  public static final String EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D =
+      "Write buffer size (%d) larger than maximum buffer size (%d)!";
+
+  public static final String EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D =
+      "Invalid compression buffer size: %d";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA =
+      "Serialized frame size is at least %d bytes, exceeding the maximum Java array size of %d bytes.";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9 =
+      "Serialized frame size is at least %d bytes, exceeding the configured Thrift max frame size of %d bytes.";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C =
+      "Serialized frame size is at least %d bytes, exceeding the maximum Java array size of %d bytes to %s.";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898 =
+      "Serialized frame size is at least %d bytes, exceeding the configured Thrift max frame size of %d bytes to %s.";
+
 }

iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java

@@ -108,4 +108,25 @@ private RpcMessages() {}
       "通过列名 %s 获取的值为 NULL。";
 
   public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 = "不支持数据类型 %s。";
+  public static final String EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA =
+      "无效的缓冲区大小：%d";
+
+  public static final String EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D =
+      "写缓冲区大小（%d）超过最大缓冲区大小（%d）！";
+
+  public static final String EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D =
+      "无效的压缩缓冲区大小：%d";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA =
+      "序列化帧大小至少为 %d 字节，超过 Java 数组最大大小 %d 字节。";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9 =
+      "序列化帧大小至少为 %d 字节，超过配置的 Thrift 最大帧大小 %d 字节。";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C =
+      "序列化帧大小至少为 %d 字节，超过 Java 数组最大大小 %d 字节，目标地址为 %s。";
+
+  public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898 =
+      "序列化帧大小至少为 %d 字节，超过配置的 Thrift 最大帧大小 %d 字节，目标地址为 %s。";
+
 }

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

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.rpc;
 
+import org.apache.iotdb.rpc.i18n.RpcMessages;
+
 import java.io.IOException;
 import java.util.Arrays;
 
@@ -35,6 +37,8 @@
  */
 class AutoResizingBuffer implements AutoCloseable {
 
+  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;
@@ -71,21 +75,38 @@ public void resizeIfNecessary(int size) throws IOException {
   }
 
   void ensureCapacity(int size) throws IOException {
+    ensureCapacity(size, MAX_ARRAY_SIZE);
+  }
+
+  void ensureCapacity(int size, int maxCapacity) throws IOException {
+    int limit = Math.min(maxCapacity, MAX_ARRAY_SIZE);
+    if (size < 0 || size > limit) {
+      throw new IllegalArgumentException(
+          String.format(RpcMessages.EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA, size));
+    }
     reserveCurrentCapacityIfReleased();
     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);
-      resize(newCapacity, size);
+      resize(Math.min(calculateNewCapacity(currentCapacity, size), limit), size);
       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;
   }
 
+  boolean isClosed() {
+    return closed;
+  }
+
   void resetToInitialCapacity() throws IOException {
     reserveCurrentCapacityIfReleased();
     resize(initialCapacity, initialCapacity);

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

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.rpc;
 
+import org.apache.iotdb.rpc.i18n.RpcMessages;
+
 import org.apache.thrift.TConfiguration;
 import org.apache.thrift.transport.TTransportException;
 
@@ -55,17 +57,48 @@ public int read(byte[] buf, int off, int len) {
 
   @Override
   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,
+          String.format(
+              RpcMessages
+                  .EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D,
+              requiredCapacity,
+              limit));
+    }
+    // A small prefix during serialization does not indicate a small frame. Only the complete
+    // frame size, supplied after flushing, may trigger shrinking.
     try {
-      // A small prefix during serialization does not indicate a small frame. Only the complete
-      // frame size, supplied after flushing, may trigger shrinking.
-      ensureCapacity(pos + len);
+      ensureCapacity((int) requiredCapacity, limit);
     } catch (IOException e) {
       throw new TTransportException(e);
     }
     System.arraycopy(toWrite, off, getBuffer(), 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 = getBuffer();
+    // A closed reusable buffer must reserve its memory again. Crossing the reuse threshold
+    // must also go through ensureCapacity so the large frame uses temporary storage.
+    if (buf.isClosed()
+        || len > array.length - pos
+        || (temporaryBuffer == null && len > maxReusableBufferSize - pos)) {
+      return false;
+    }
+    System.arraycopy(toWrite, off, array, pos, len);
+    pos += len;
+    return true;
+  }
+
   public int getPos() {
     return pos;
   }
@@ -101,20 +134,26 @@ public void resizeIfNecessary(int size) throws IOException {
   }
 
   private void ensureCapacity(int size) throws IOException {
+    ensureCapacity(size, AutoResizingBuffer.MAX_ARRAY_SIZE);
+  }
+
+  private void ensureCapacity(int size, int maxCapacity) throws IOException {
     if (size > maxReusableBufferSize) {
       if (temporaryBuffer == null || temporaryBuffer.length < size) {
         // Preserve an unfinished frame when incremental serialization crosses the threshold.
         // Keep amortized growth for subsequent small writes within the same temporary frame.
         int capacity =
             temporaryBuffer == null
                 ? size
-                : Math.max(size, temporaryBuffer.length + (temporaryBuffer.length >> 1));
+                : Math.min(
+                    AutoResizingBuffer.calculateNewCapacity(temporaryBuffer.length, size),
+                    maxCapacity);
         byte[] newBuffer = Arrays.copyOf(getBuffer(), capacity);
         buf.resetToInitialCapacity();
         temporaryBuffer = newBuffer;
       }
     } else {
-      buf.ensureCapacity(size);
+      buf.ensureCapacity(size, maxCapacity);
     }
   }
 

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

@@ -19,6 +19,8 @@
 
 package org.apache.iotdb.rpc;
 
+import org.apache.iotdb.rpc.i18n.RpcMessages;
+
 import org.apache.thrift.transport.TTransport;
 import org.apache.thrift.transport.TTransportException;
 import org.apache.thrift.transport.layered.TFramedTransport;
@@ -75,6 +77,13 @@ protected void closeAllocatedBuffers() {
     }
   }
 
+  @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 {
     // Discard the previous frame's budget before reading the next frame header and payload.
@@ -88,6 +97,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);
@@ -105,6 +115,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);
@@ -128,6 +139,17 @@ 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,
+          String.format(RpcMessages.EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D, 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

@@ -43,9 +43,8 @@
   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;
 
@@ -106,7 +105,11 @@
       readBuffer =
           new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize);
       writeBuffer =
-          new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize);
+          new AutoScalingBufferWriteTransport(
+              Math.min(
+                  thriftDefaultBufferSize,
+                  Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE)),
+              maxReusableBufferSize);
     } catch (IOException e) {
       closeAllocatedBuffers();
       throw new TTransportException(e);
@@ -383,7 +386,50 @@
 
   @Override
   public void write(byte[] buf, int off, int len) throws TTransportException {
-    writeBuffer.write(buf, off, len);
+    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.
+      // TSocket.close() clears its socket, so retain the destination before invalidating it.
+      SocketAddress remoteAddress = null;
+      if (underlying instanceof TSocket socket && socket.getSocket() != null) {
+        remoteAddress = socket.getSocket().getRemoteSocketAddress();
+      }
+      close();
+      String message;
+      if (remoteAddress == null) {
+        message =
+            String.format(
+                maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE
+                    ? RpcMessages
+                        .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA
+                    : RpcMessages
+                        .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9,
+                size,
+                maxSize);
+      } else {
+        message =
+            String.format(
+                maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE
+                    ? RpcMessages
+                        .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C
+                    : RpcMessages
+                        .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898,
+                size,
+                maxSize,
+                remoteAddress);
+      }
+      throw new TTransportException(TTransportException.CORRUPTED_DATA, message);
+    }
+    writeBuffer.write(buf, off, len, maxSize);
+  }
+
+  protected int getMaxWriteBufferSize() {
+    return Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE);
   }
 
   public TTransport getSocket() {