[To dev/1.3] Fix RPC buffer resizing and duplicate flush submissions (#18918)

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

@@ -22,14 +22,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. 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 {
 
@@ -44,22 +45,27 @@ public AutoResizingBuffer(int initialCapacity) {
   }
 
   public void resizeIfNecessary(int size) {
+    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) {
+      array = Arrays.copyOf(array, size + (currentCapacity - size) / 2);
+      bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
+      lastShrinkTime = System.currentTimeMillis();
+    }
+  }
+
+  void ensureCapacity(int 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);
       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
-      array = Arrays.copyOf(array, size + (currentCapacity - size) / 2);
-      bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME;
-      lastShrinkTime = System.currentTimeMillis();
     }
   }
 

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

@@ -44,7 +44,9 @@ public int read(byte[] buf, int off, int len) {
 
   @Override
   public void write(byte[] toWrite, int off, int len) {
-    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);
     System.arraycopy(toWrite, off, buf.array(), pos, len);
     pos += len;
   }

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

@@ -79,9 +79,7 @@ public void flush() throws TTransportException {
     }
 
     writeBuffer.reset();
-    if (thriftDefaultBufferSize < length) {
-      writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
-    }
+    writeBuffer.resizeIfNecessary(length);
     underlying.flush();
   }
 

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

@@ -260,9 +260,7 @@ public void flush() throws TTransportException {
     underlying.write(i32buf, 0, 4);
     underlying.write(writeBuffer.getBuffer(), 0, length);
     writeBuffer.reset();
-    if (length > thriftDefaultBufferSize) {
-      writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
-    }
+    writeBuffer.resizeIfNecessary(length);
     underlying.flush();
   }
 

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

@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.rpc;
+
+import org.apache.thrift.transport.TMemoryBuffer;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.util.Random;
+
+public class AutoResizingBufferTest {
+
+  @Test
+  public void testGrowthPreservesDataAndHeadroom() {
+    AutoResizingBuffer buffer = new AutoResizingBuffer(100);
+    buffer.array()[99] = 42;
+    buffer.resizeIfNecessary(110);
+    Assert.assertEquals(150, buffer.array().length);
+    Assert.assertEquals(42, buffer.array()[99]);
+
+    buffer.resizeIfNecessary(1000);
+    Assert.assertEquals(1000, buffer.array().length);
+    Assert.assertEquals(42, buffer.array()[99]);
+  }
+
+  @Test
+  public void testShrinkForRequestsAtOrBelowInitialCapacity() throws Exception {
+    for (int requestSize : new int[] {0, 1, 100}) {
+      AutoResizingBuffer buffer = new AutoResizingBuffer(100);
+      buffer.resizeIfNecessary(1000);
+      buffer.array()[99] = 42;
+      for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+        buffer.resizeIfNecessary(requestSize);
+      }
+
+      Assert.assertEquals(550, buffer.array().length);
+      Assert.assertEquals(42, buffer.array()[99]);
+      for (int i = 0; i < 20; i++) {
+        setLastShrinkTime(buffer, 0);
+        for (int j = 0; j <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; j++) {
+          buffer.resizeIfNecessary(requestSize);
+        }
+        Assert.assertTrue(buffer.array().length >= 100);
+      }
+    }
+  }
+
+  @Test
+  public void testSmallFramesRespectShrinkCountAndCooldown() throws Exception {
+    AutoResizingBuffer buffer = new AutoResizingBuffer(100);
+    buffer.resizeIfNecessary(1000);
+    byte[] largeBuffer = buffer.array();
+    setLastShrinkTime(buffer, System.currentTimeMillis() + RpcUtils.MIN_SHRINK_INTERVAL);
+
+    for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+      buffer.resizeIfNecessary(101);
+    }
+    Assert.assertSame(largeBuffer, buffer.array());
+
+    setLastShrinkTime(buffer, 0);
+    buffer.resizeIfNecessary(101);
+    Assert.assertEquals(550, buffer.array().length);
+    byte[] shrunkBuffer = buffer.array();
+    for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+      buffer.resizeIfNecessary(101);
+    }
+    Assert.assertSame(shrunkBuffer, buffer.array());
+
+    AutoResizingBuffer countLimitedBuffer = new AutoResizingBuffer(100);
+    countLimitedBuffer.resizeIfNecessary(1000);
+    byte[] originalBuffer = countLimitedBuffer.array();
+    for (int i = 0; i < RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+      countLimitedBuffer.resizeIfNecessary(101);
+      Assert.assertSame(originalBuffer, countLimitedBuffer.array());
+    }
+    countLimitedBuffer.resizeIfNecessary(101);
+    Assert.assertEquals(550, countLimitedBuffer.array().length);
+  }
+
+  @Test
+  public void testFramedTransportsShrinkAfterLargeFrames() throws Exception {
+    for (boolean compressed : new boolean[] {false, true}) {
+      try (TElasticFramedTransport transport = newFramedTransport(compressed)) {
+        byte[] payload = new byte[1000];
+        payload[999] = 42;
+        transport.write(payload);
+        transport.flush();
+        byte[] actual = new byte[payload.length];
+        transport.readAll(actual, 0, actual.length);
+        Assert.assertArrayEquals(payload, actual);
+        byte[] largeWriteBuffer = transport.writeBuffer.getBuffer();
+        byte[] largeReadBuffer = transport.readBuffer.getBuffer();
+        long cooldown = System.currentTimeMillis() + RpcUtils.MIN_SHRINK_INTERVAL;
+        setLastShrinkTime(transport.writeBuffer, cooldown);
+        setLastShrinkTime(transport.readBuffer, cooldown);
+
+        for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+          transport.write(new byte[] {42});
+          transport.flush();
+          byte[] smallFrame = new byte[1];
+          transport.readAll(smallFrame, 0, smallFrame.length);
+          Assert.assertEquals(42, smallFrame[0]);
+        }
+        Assert.assertSame(largeWriteBuffer, transport.writeBuffer.getBuffer());
+        Assert.assertSame(largeReadBuffer, transport.readBuffer.getBuffer());
+
+        setLastShrinkTime(transport.writeBuffer, 0);
+        setLastShrinkTime(transport.readBuffer, 0);
+        transport.write(new byte[] {42});
+        transport.flush();
+        byte[] smallFrame = new byte[1];
+        transport.readAll(smallFrame, 0, 1);
+        Assert.assertEquals(42, smallFrame[0]);
+        Assert.assertTrue(transport.writeBuffer.getBuffer().length < payload.length);
+        Assert.assertTrue(transport.readBuffer.getBuffer().length < payload.length);
+        Assert.assertTrue(transport.writeBuffer.getBuffer().length >= 100);
+        Assert.assertTrue(transport.readBuffer.getBuffer().length >= 100);
+
+        byte[] shrunkBuffer = transport.writeBuffer.getBuffer();
+        for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+          transport.write(new byte[] {42});
+          transport.flush();
+          transport.readAll(smallFrame, 0, 1);
+          Assert.assertEquals(42, smallFrame[0]);
+        }
+        Assert.assertSame(shrunkBuffer, transport.writeBuffer.getBuffer());
+      }
+    }
+  }
+
+  @Test
+  public void testRepeatedLargeFramesKeepWriteCapacity() throws Exception {
+    for (boolean compressed : new boolean[] {false, true}) {
+      try (TElasticFramedTransport transport = newFramedTransport(compressed)) {
+        byte[] payload = new byte[1000];
+        new Random(0).nextBytes(payload);
+        transport.write(payload);
+        transport.flush();
+        byte[] actual = new byte[payload.length];
+        transport.readAll(actual, 0, actual.length);
+        Assert.assertArrayEquals(payload, actual);
+        byte[] writeBuffer = transport.writeBuffer.getBuffer();
+        for (int i = 0; i < 3 * (RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1); i++) {
+          setLastShrinkTime(transport.writeBuffer, 0);
+          transport.write(payload, 0, 200);
+          transport.write(payload, 200, payload.length - 200);
+          transport.flush();
+          Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+          transport.readAll(actual, 0, actual.length);
+          Assert.assertArrayEquals(payload, actual);
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testSmallWritesDoNotShrinkAnUnfinishedFrame() throws Exception {
+    for (boolean compressed : new boolean[] {false, true}) {
+      try (TElasticFramedTransport transport = newFramedTransport(compressed)) {
+        byte[] payload = new byte[1000];
+        new Random(0).nextBytes(payload);
+        transport.write(payload);
+        transport.flush();
+        byte[] actual = new byte[payload.length];
+        transport.readAll(actual, 0, actual.length);
+        Assert.assertArrayEquals(payload, actual);
+        byte[] writeBuffer = transport.writeBuffer.getBuffer();
+        setLastShrinkTime(transport.writeBuffer, 0);
+
+        int fragmentLength = 32;
+        int prefixLength = fragmentLength * (RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 5);
+        for (int offset = 0; offset < prefixLength; offset += fragmentLength) {
+          transport.write(payload, offset, fragmentLength);
+          Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+          Assert.assertEquals(offset + fragmentLength, transport.writeBuffer.getPos());
+        }
+        transport.write(payload, prefixLength, payload.length - prefixLength);
+        Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+        transport.flush();
+        transport.readAll(actual, 0, actual.length);
+        Assert.assertArrayEquals(payload, actual);
+      }
+    }
+  }
+
+  private void setLastShrinkTime(AutoResizingBuffer buffer, long lastShrinkTime) throws Exception {
+    Field field = AutoResizingBuffer.class.getDeclaredField("lastShrinkTime");
+    field.setAccessible(true);
+    field.set(buffer, lastShrinkTime);
+  }
+
+  private void setLastShrinkTime(Object transport, long lastShrinkTime) throws Exception {
+    Field field = transport.getClass().getDeclaredField("buf");
+    field.setAccessible(true);
+    setLastShrinkTime((AutoResizingBuffer) field.get(transport), lastShrinkTime);
+  }
+
+  private TElasticFramedTransport newFramedTransport(boolean compressed) throws Exception {
+    return compressed
+        ? new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 100, 2048, true)
+        : new TElasticFramedTransport(new TMemoryBuffer(0), 100, 2048, true);
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfo.java

@@ -493,8 +493,12 @@ private boolean chooseMemTablesToMarkFlush(TsFileProcessor currentTsFileProcesso
       }
       TsFileProcessor selectedTsFileProcessor = allTsFileProcessors.peek();
       memCost += selectedTsFileProcessor.getWorkMemTableRamCost();
-      selectedTsFileProcessor.setWorkMemTableShouldFlush();
-      flushTaskSubmitThreadPool.submit(selectedTsFileProcessor::submitAFlushTask);
+      // A marked memtable already has a flush pending. Count its memory towards the target,
+      // but do not enqueue another task on every report while the region write lock is held.
+      if (!selectedTsFileProcessor.getWorkMemTable().shouldFlush()) {
+        selectedTsFileProcessor.setWorkMemTableShouldFlush();
+        flushTaskSubmitThreadPool.submit(selectedTsFileProcessor::submitAFlushTask);
+      }
       if (selectedTsFileProcessor == currentTsFileProcessor) {
         isCurrentTsFileProcessorSelected = true;
       }