Fix receiver write admission stalls and RPC buffer allocation failures

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

@@ -23,14 +23,15 @@
 import java.util.Arrays;
 
 /**
- * Helper class that wraps a byte[] so that it can expand and be reused. Users should call
- * resizeIfNecessary to make sure the buffer has suitable capacity, and then use the array as
- * needed.
+ * Helper class that wraps a byte[] so that it can expand and be reused. Incremental writes use
+ * ensureCapacity; resizeIfNecessary also considers shrinking and must receive a complete frame's
+ * size.
  *
  * <p>Resizing policies: Expanding: If the required size > current capacity * 1.5, expand to the
- * required size, otherwise expand to current capacity * 1.5. Shrinking: If initial size < the
- * required size < current capacity * 0.6, and such small requests last for more than 5 times,
- * shrink to the middle of the required size and current capacity.
+ * required size, otherwise expand to current capacity * 1.5 when the memory budget allows it.
+ * Shrinking: If the required size < current capacity * 0.6, and such small requests last for more
+ * than 5 times, shrink to the middle of the required size and current capacity. The required size
+ * is bounded below by the initial capacity.
  */
 class AutoResizingBuffer implements AutoCloseable {
 
@@ -56,23 +57,28 @@ public AutoResizingBuffer(int initialCapacity) throws IOException {
   }
 
   public void resizeIfNecessary(int size) throws IOException {
-    reserveCurrentCapacityIfReleased();
+    ensureCapacity(size);
+    size = Math.max(size, initialCapacity);
     final int currentCapacity = this.array.length;
     final double loadFactor = 0.6;
+    if (currentCapacity * loadFactor > size
+        && bufTooLargeCounter-- <= 0
+        && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) {
+      resize(size + (currentCapacity - size) / 2, size);
+      bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
+      lastShrinkTime = System.currentTimeMillis();
+    }
+  }
+
+  void ensureCapacity(int size) throws IOException {
+    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);
+      resize(newCapacity, size);
       bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
-    } else if (size > initialCapacity
-        && currentCapacity * loadFactor > size
-        && bufTooLargeCounter-- <= 0
-        && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) {
-      // do not resize if it is reading the request size and do not shrink too often
-      resize(size + (currentCapacity - size) / 2);
-      bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
-      lastShrinkTime = System.currentTimeMillis();
     }
   }
 
@@ -89,11 +95,14 @@ public void close() {
     }
   }
 
-  private void resize(int newCapacity) throws IOException {
+  private void resize(int newCapacity, int requiredCapacity) throws IOException {
     final int currentCapacity = array.length;
     if (newCapacity > currentCapacity) {
-      final int delta = newCapacity - currentCapacity;
-      memoryAccount.allocate(delta);
+      final int delta =
+          (int)
+              memoryAccount.allocate(
+                  newCapacity - currentCapacity, requiredCapacity - currentCapacity);
+      newCapacity = currentCapacity + delta;
       try {
         array = Arrays.copyOf(array, newCapacity);
         accountedCapacity += delta;

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

@@ -124,6 +124,18 @@ void allocate(long sizeInBytes) throws IOException {
       AutoResizingBufferMemoryManager.allocate(memoryControl, sizeInBytes);
     }
 
+    long allocate(long preferredSizeInBytes, long requiredSizeInBytes) throws IOException {
+      // Growth headroom is optional. If it does not fit, only wait for the bytes actually needed
+      // by the frame instead of rejecting a request that still fits within the memory budget.
+      if (preferredSizeInBytes > requiredSizeInBytes
+          && memoryControl.allocate(preferredSizeInBytes)) {
+        memoryAllocationCount.incrementAndGet();
+        return preferredSizeInBytes;
+      }
+      allocate(requiredSizeInBytes);
+      return requiredSizeInBytes;
+    }
+
     void release(long sizeInBytes) {
       AutoResizingBufferMemoryManager.release(memoryControl, sizeInBytes);
     }

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

@@ -47,7 +47,9 @@ public int read(byte[] buf, int off, int len) {
   @Override
   public void write(byte[] toWrite, int off, int len) throws TTransportException {
     try {
-      buf.resizeIfNecessary(pos + len);
+      // 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);
     } catch (IOException e) {
       throw new TTransportException(e);
     }

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

@@ -100,12 +100,10 @@ public void flush() throws TTransportException {
     }
 
     writeBuffer.reset();
-    if (thriftDefaultBufferSize < length) {
-      try {
-        writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
-      } catch (IOException e) {
-        throw new TTransportException(e);
-      }
+    try {
+      writeBuffer.resizeIfNecessary(length);
+    } catch (IOException e) {
+      throw new TTransportException(e);
     }
     underlying.flush();
   }

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

@@ -317,12 +317,10 @@ public void flush() throws TTransportException {
     underlying.write(i32buf, 0, 4);
     underlying.write(writeBuffer.getBuffer(), 0, length);
     writeBuffer.reset();
-    if (length > thriftDefaultBufferSize) {
-      try {
-        writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
-      } catch (IOException e) {
-        throw new TTransportException(e);
-      }
+    try {
+      writeBuffer.resizeIfNecessary(length);
+    } catch (IOException e) {
+      throw new TTransportException(e);
     }
     underlying.flush();
   }