Add configurable Thrift buffer reuse thresholds to Java clients

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

@@ -86,6 +86,13 @@ public byte[] array() {
     return this.array;
   }
 
+  void resetToInitialCapacity() throws IOException {
+    reserveCurrentCapacityIfReleased();
+    resize(initialCapacity, initialCapacity);
+    bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
+    lastShrinkTime = System.currentTimeMillis();
+  }
+
   @Override
   public void close() {
     if (!closed) {

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

@@ -28,28 +28,36 @@
 public class AutoScalingBufferReadTransport extends NonOpenTransport {
 
   private final AutoResizingBuffer buf;
+  private final int maxReusableBufferSize;
+  private byte[] temporaryBuffer;
   private int pos = 0;
   private int limit = 0;
 
   public AutoScalingBufferReadTransport(int initialCapacity) throws IOException {
+    this(initialCapacity, Integer.MAX_VALUE);
+  }
+
+  AutoScalingBufferReadTransport(int initialCapacity, int maxReusableBufferSize)
+      throws IOException {
     this.buf = new AutoResizingBuffer(initialCapacity);
+    this.maxReusableBufferSize = maxReusableBufferSize;
   }
 
   public void fill(TTransport inTrans, int length) throws TTransportException {
     try {
-      buf.resizeIfNecessary(length);
+      resizeIfNecessary(length);
     } catch (IOException e) {
       throw new TTransportException(e);
     }
-    inTrans.readAll(buf.array(), 0, length);
+    inTrans.readAll(getBuffer(), 0, length);
     pos = 0;
     limit = length;
   }
 
   @Override
   public final int read(byte[] target, int off, int len) {
     int amtToRead = Math.min(len, getBytesRemainingInBuffer());
-    System.arraycopy(buf.array(), pos, target, off, amtToRead);
+    System.arraycopy(getBuffer(), pos, target, off, amtToRead);
     consumeBuffer(amtToRead);
     return amtToRead;
   }
@@ -62,6 +70,12 @@ public void write(byte[] buf, int off, int len) {
   @Override
   public final void consumeBuffer(int len) {
     pos += len;
+    if (pos == limit && temporaryBuffer != null) {
+      // Protocols may keep their own reference to a consumed binary value. Drop only our reference.
+      temporaryBuffer = null;
+      pos = 0;
+      limit = 0;
+    }
   }
 
   @Override
@@ -82,7 +96,7 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {
 
   @Override
   public final byte[] getBuffer() {
-    return buf.array();
+    return temporaryBuffer == null ? buf.array() : temporaryBuffer;
   }
 
   @Override
@@ -96,12 +110,22 @@ public final int getBytesRemainingInBuffer() {
   }
 
   public void resizeIfNecessary(int size) throws IOException {
-    buf.resizeIfNecessary(size);
+    if (size > maxReusableBufferSize) {
+      byte[] newBuffer = new byte[size];
+      buf.resetToInitialCapacity();
+      temporaryBuffer = newBuffer;
+    } else {
+      buf.resizeIfNecessary(size);
+      temporaryBuffer = null;
+    }
   }
 
   @Override
   public void close() {
     super.close();
+    temporaryBuffer = null;
+    pos = 0;
+    limit = 0;
     buf.close();
   }
 

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

@@ -23,6 +23,7 @@
 import org.apache.thrift.transport.TTransportException;
 
 import java.io.IOException;
+import java.util.Arrays;
 
 /**
  * Note that this class is mainly copied from class {@link
@@ -32,10 +33,18 @@
 public class AutoScalingBufferWriteTransport extends NonOpenTransport {
 
   private final AutoResizingBuffer buf;
+  private final int maxReusableBufferSize;
+  private byte[] temporaryBuffer;
   private int pos;
 
   public AutoScalingBufferWriteTransport(int initialCapacity) throws IOException {
+    this(initialCapacity, Integer.MAX_VALUE);
+  }
+
+  AutoScalingBufferWriteTransport(int initialCapacity, int maxReusableBufferSize)
+      throws IOException {
     this.buf = new AutoResizingBuffer(initialCapacity);
+    this.maxReusableBufferSize = maxReusableBufferSize;
     this.pos = 0;
   }
 
@@ -49,11 +58,11 @@ public void write(byte[] toWrite, int off, int len) throws TTransportException {
     try {
       // 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);
+      ensureCapacity(pos + len);
     } catch (IOException e) {
       throw new TTransportException(e);
     }
-    System.arraycopy(toWrite, off, buf.array(), pos, len);
+    System.arraycopy(toWrite, off, getBuffer(), pos, len);
     pos += len;
   }
 
@@ -63,21 +72,55 @@ public int getPos() {
 
   public void reset() {
     pos = 0;
+    temporaryBuffer = null;
+  }
+
+  void resetAndResize() throws IOException {
+    int length = pos;
+    reset();
+    // A completed large frame must not recreate the temporary array just released by reset().
+    if (length <= maxReusableBufferSize) {
+      buf.resizeIfNecessary(length);
+    }
   }
 
   @Override
   public void close() {
     super.close();
+    reset();
     buf.close();
   }
 
   public void resizeIfNecessary(int size) throws IOException {
-    buf.resizeIfNecessary(size);
+    if (size > maxReusableBufferSize) {
+      ensureCapacity(size);
+    } else {
+      buf.resizeIfNecessary(size);
+      temporaryBuffer = null;
+    }
+  }
+
+  private void ensureCapacity(int size) 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));
+        byte[] newBuffer = Arrays.copyOf(getBuffer(), capacity);
+        buf.resetToInitialCapacity();
+        temporaryBuffer = newBuffer;
+      }
+    } else {
+      buf.ensureCapacity(size);
+    }
   }
 
   @Override
   public byte[] getBuffer() {
-    return buf.array();
+    return temporaryBuffer == null ? buf.array() : temporaryBuffer;
   }
 
   @Override

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

@@ -42,14 +42,28 @@ public static DeepCopyRpcTransportFactory getInstance(
     return create(USE_SNAPPY, thriftDefaultBufferSize, thriftMaxFrameSize);
   }
 
+  /** Creates a client factory whose reusable-buffer threshold is independent of other clients. */
+  public static DeepCopyRpcTransportFactory getInstance(
+      int thriftDefaultBufferSize, int thriftMaxFrameSize, int maxReusableBufferSize) {
+    return create(USE_SNAPPY, thriftDefaultBufferSize, thriftMaxFrameSize, maxReusableBufferSize);
+  }
+
   private static DeepCopyRpcTransportFactory create(
       boolean useSnappy, int thriftDefaultBufferSize, int thriftMaxFrameSize) {
+    return create(useSnappy, thriftDefaultBufferSize, thriftMaxFrameSize, Integer.MAX_VALUE);
+  }
+
+  private static DeepCopyRpcTransportFactory create(
+      boolean useSnappy,
+      int thriftDefaultBufferSize,
+      int thriftMaxFrameSize,
+      int maxReusableBufferSize) {
     return useSnappy
         ? new DeepCopyRpcTransportFactory(
             new TimeoutChangeableTSnappyFramedTransport.Factory(
-                thriftDefaultBufferSize, thriftMaxFrameSize, true))
+                thriftDefaultBufferSize, thriftMaxFrameSize, true, maxReusableBufferSize))
         : new DeepCopyRpcTransportFactory(
             new TimeoutChangeableTFastFramedTransport.Factory(
-                thriftDefaultBufferSize, thriftMaxFrameSize, true));
+                thriftDefaultBufferSize, thriftMaxFrameSize, true, maxReusableBufferSize));
   }
 }

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

@@ -36,10 +36,28 @@ protected TCompressedElasticFramedTransport(
       int thriftMaxFrameSize,
       boolean copyBinary)
       throws TTransportException {
-    super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary);
+    this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE);
+  }
+
+  /**
+   * @param maxReusableBufferSize requests larger than this number of bytes use temporary arrays in
+   *     all four buffers and reset their reusable storage to {@code thriftDefaultBufferSize};
+   *     {@link Integer#MAX_VALUE} preserves the original resizing policy
+   */
+  protected TCompressedElasticFramedTransport(
+      TTransport underlying,
+      int thriftDefaultBufferSize,
+      int thriftMaxFrameSize,
+      boolean copyBinary,
+      int maxReusableBufferSize)
+      throws TTransportException {
+    super(
+        underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize);
     try {
-      writeCompressBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize);
-      readCompressBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize);
+      writeCompressBuffer =
+          new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize);
+      readCompressBuffer =
+          new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize);
     } catch (IOException e) {
       closeAllocatedBuffers();
       throw new TTransportException(e);
@@ -97,11 +115,13 @@ public void flush() throws TTransportException {
       underlying.write(writeCompressBuffer.getBuffer(), 0, compressedLength);
     } catch (IOException e) {
       throw new TTransportException(e);
+    } finally {
+      // The compressed output has been handed to the underlying transport, so do not retain it.
+      writeCompressBuffer.reset();
     }
 
-    writeBuffer.reset();
     try {
-      writeBuffer.resizeIfNecessary(length);
+      writeBuffer.resetAndResize();
     } catch (IOException e) {
       throw new TTransportException(e);
     }

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

@@ -87,14 +87,26 @@ public TElasticFramedTransport(
       int thriftMaxFrameSize,
       boolean copyBinary)
       throws TTransportException {
+    this(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE);
+  }
+
+  protected TElasticFramedTransport(
+      TTransport underlying,
+      int thriftDefaultBufferSize,
+      int thriftMaxFrameSize,
+      boolean copyBinary,
+      int maxReusableBufferSize)
+      throws TTransportException {
     this.underlying = underlying;
     this.thriftDefaultBufferSize = thriftDefaultBufferSize;
     this.thriftMaxFrameSize = thriftMaxFrameSize;
     this.copyBinary = copyBinary;
     alignUnderlyingConfiguration();
     try {
-      readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize);
-      writeBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize);
+      readBuffer =
+          new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize);
+      writeBuffer =
+          new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize);
     } catch (IOException e) {
       closeAllocatedBuffers();
       throw new TTransportException(e);
@@ -316,9 +328,8 @@ public void flush() throws TTransportException {
     TFramedTransport.encodeFrameSize(length, i32buf);
     underlying.write(i32buf, 0, 4);
     underlying.write(writeBuffer.getBuffer(), 0, length);
-    writeBuffer.reset();
     try {
-      writeBuffer.resizeIfNecessary(length);
+      writeBuffer.resetAndResize();
     } catch (IOException e) {
       throw new TTransportException(e);
     }

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

@@ -29,6 +29,8 @@ public class TSnappyElasticFramedTransport extends TCompressedElasticFramedTrans
 
   public static class Factory extends TElasticFramedTransport.Factory {
 
+    private final int maxReusableBufferSize;
+
     public Factory() {
       this(RpcUtils.THRIFT_DEFAULT_BUF_CAPACITY, RpcUtils.THRIFT_FRAME_MAX_SIZE, true);
     }
@@ -38,13 +40,25 @@ public Factory(int thriftDefaultBufferSize, boolean copyBinary) {
     }
 
     public Factory(int thriftDefaultBufferSize, int thriftMaxFrameSize, boolean copyBinary) {
+      this(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, Integer.MAX_VALUE);
+    }
+
+    /**
+     * @param maxReusableBufferSize buffer requests above this size in bytes use temporary arrays
+     */
+    public Factory(
+        int thriftDefaultBufferSize,
+        int thriftMaxFrameSize,
+        boolean copyBinary,
+        int maxReusableBufferSize) {
       super(thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary);
+      this.maxReusableBufferSize = maxReusableBufferSize;
     }
 
     @Override
     public TTransport getTransport(TTransport trans) throws TTransportException {
       return new TSnappyElasticFramedTransport(
-          trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary);
+          trans, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize);
     }
   }
 
@@ -61,6 +75,20 @@ public TSnappyElasticFramedTransport(
     super(underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary);
   }
 
+  /**
+   * @param maxReusableBufferSize buffer requests above this size in bytes use temporary arrays
+   */
+  public TSnappyElasticFramedTransport(
+      TTransport underlying,
+      int thriftDefaultBufferSize,
+      int thriftMaxFrameSize,
+      boolean copyBinary,
+      int maxReusableBufferSize)
+      throws TTransportException {
+    super(
+        underlying, thriftDefaultBufferSize, thriftMaxFrameSize, copyBinary, maxReusableBufferSize);
+  }
+
   @Override
   protected int uncompressedLength(byte[] buf, int off, int len) throws IOException {
     return Snappy.uncompressedLength(buf, off, len);