This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 83d419796fb Fix receiver write admission stalls and RPC buffer
allocation failures (#18918)
83d419796fb is described below
commit 83d419796fb91a6cccb874414820d3865ab34b98
Author: Caideyipi <[email protected]>
AuthorDate: Wed Oct 21 12:21:08 2026 +0800
Fix receiver write admission stalls and RPC buffer allocation failures
(#18918)
* Fix receiver write admission and RPC buffer memory pressure
* Preserve RPC frame buffers and avoid duplicate flush submissions
* Revert fallback flush trigger in memory admission
---
.../org/apache/iotdb/rpc/AutoResizingBuffer.java | 47 +--
.../iotdb/rpc/AutoResizingBufferMemoryManager.java | 12 +
.../iotdb/rpc/AutoScalingBufferWriteTransport.java | 4 +-
.../rpc/TCompressedElasticFramedTransport.java | 10 +-
.../apache/iotdb/rpc/TElasticFramedTransport.java | 10 +-
.../apache/iotdb/rpc/AutoResizingBufferTest.java | 347 +++++++++++++++++++++
.../db/storageengine/rescon/memory/SystemInfo.java | 8 +-
.../rescon/memory/SystemInfoTest.java | 232 ++++++++++++++
8 files changed, 636 insertions(+), 34 deletions(-)
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java
index aa0ea40541c..d9772985dc6 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java
@@ -23,14 +23,15 @@ import java.io.IOException;
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 @@ class AutoResizingBuffer implements AutoCloseable {
}
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 @@ class AutoResizingBuffer implements AutoCloseable {
}
}
- 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;
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryManager.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryManager.java
index d33a15b44cb..bbdbc8bda06 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryManager.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryManager.java
@@ -124,6 +124,18 @@ public final class AutoResizingBufferMemoryManager {
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);
}
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java
index 24278067fcd..eab45f3da48 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java
@@ -47,7 +47,9 @@ public class AutoScalingBufferWriteTransport extends
NonOpenTransport {
@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);
}
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java
index b4c6c0994f5..36cca36b9b1 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TCompressedElasticFramedTransport.java
@@ -100,12 +100,10 @@ public abstract class TCompressedElasticFramedTransport
extends TElasticFramedTr
}
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();
}
diff --git
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java
index 339dee595e7..e15c9f59424 100644
---
a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java
+++
b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java
@@ -317,12 +317,10 @@ public class TElasticFramedTransport extends TTransport {
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();
}
diff --git
a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/AutoResizingBufferTest.java
b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/AutoResizingBufferTest.java
index 164d0f9e287..18aed44eec3 100644
---
a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/AutoResizingBufferTest.java
+++
b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/AutoResizingBufferTest.java
@@ -27,6 +27,8 @@ import org.junit.Test;
import java.io.IOException;
import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Random;
import java.util.concurrent.atomic.AtomicLong;
public class AutoResizingBufferTest {
@@ -60,6 +62,336 @@ public class AutoResizingBufferTest {
Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
}
+ @Test
+ public void testGrowToRequiredCapacityWhenPreferredCapacityExceedsBudget()
throws Exception {
+ memoryControl.setLimit(120);
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+ AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
+
+ try (AutoResizingBuffer buffer = new AutoResizingBuffer(100)) {
+ buffer.array()[99] = 42;
+ buffer.resizeIfNecessary(110);
+
+ Assert.assertEquals(110, buffer.array().length);
+ Assert.assertEquals(110, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(42, buffer.array()[99]);
+ Assert.assertEquals(2,
AutoResizingBufferMemoryManager.getMemoryAllocationCount());
+ Assert.assertEquals(0,
AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount());
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+
+ @Test
+ public void testKeepPreferredGrowthWhenBudgetAllows() throws Exception {
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+
+ try (AutoResizingBuffer buffer = new AutoResizingBuffer(100)) {
+ buffer.resizeIfNecessary(110);
+ Assert.assertEquals(150, buffer.array().length);
+ Assert.assertEquals(150, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(2,
AutoResizingBufferMemoryManager.getMemoryAllocationCount());
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+
+ @Test
+ public void testShrinkForRequestsAtOrBelowInitialCapacity() throws Exception
{
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+
+ for (int requestSize : new int[] {0, 1, 100}) {
+ try (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(550, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(42, buffer.array()[99]);
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ }
+
+ @Test
+ public void testSmallFramesReleaseCapacityForAnotherConnection() throws
Exception {
+ memoryControl.setLimit(800);
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+ AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
+
+ AutoScalingBufferReadTransport first = new
AutoScalingBufferReadTransport(100);
+ AutoScalingBufferReadTransport second = new
AutoScalingBufferReadTransport(100);
+ try {
+ TMemoryBuffer input = new TMemoryBuffer(600);
+ byte[] payload = new byte[600];
+ payload[0] = 42;
+ input.write(payload);
+ first.fill(input, payload.length);
+ Assert.assertEquals(700, memoryControl.getUsedMemoryInBytes());
+ first.consumeBuffer(payload.length);
+
+ for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) {
+ input.write(new byte[] {42});
+ first.fill(input, 1);
+ first.consumeBuffer(1);
+ }
+ Assert.assertEquals(450, memoryControl.getUsedMemoryInBytes());
+
+ input.write(payload, 0, 400);
+ second.fill(input, 400);
+ Assert.assertEquals(750, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(42, first.getBuffer()[0]);
+ Assert.assertEquals(42, second.getBuffer()[0]);
+ Assert.assertEquals(400, second.getBytesRemainingInBuffer());
+ } finally {
+ first.close();
+ second.close();
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+
+ @Test
+ public void testFramedTransportsShrinkAfterLargeFrames() throws Exception {
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+
+ for (boolean compressed : new boolean[] {false, true}) {
+ TMemoryBuffer underlying = new TMemoryBuffer(0);
+ try (TElasticFramedTransport transport =
+ compressed
+ ? new TSnappyElasticFramedTransport(underlying, 100, 2048, true)
+ : new TElasticFramedTransport(underlying, 100, 2048, true)) {
+ 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);
+ long largeFrameMemory = memoryControl.getUsedMemoryInBytes();
+ 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);
+ Assert.assertTrue(memoryControl.getUsedMemoryInBytes() <
largeFrameMemory);
+
+ 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());
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ }
+
+ @Test
+ public void testRepeatedLargeFramesKeepWriteCapacity() throws Exception {
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+
+ 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);
+ long allocationCount =
AutoResizingBufferMemoryManager.getMemoryAllocationCount();
+ transport.write(payload);
+ transport.flush();
+ Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+ Assert.assertEquals(
+ allocationCount,
AutoResizingBufferMemoryManager.getMemoryAllocationCount());
+ transport.readAll(actual, 0, actual.length);
+ Assert.assertArrayEquals(payload, actual);
+ }
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ }
+
+ @Test
+ public void testSmallWritesDoNotShrinkAnUnfinishedFrame() throws Exception {
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+
+ 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 prefixLength = RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1;
+ for (int i = 0; i < prefixLength; i++) {
+ transport.write(payload, i, 1);
+ Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+ Assert.assertEquals(i + 1, 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);
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ }
+
+ @Test
+ public void testLargeFrameReuseDoesNotStarveAnotherConnection() throws
Exception {
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+ AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
+
+ for (boolean compressed : new boolean[] {false, true}) {
+ memoryControl.setLimit(Long.MAX_VALUE);
+ try (TElasticFramedTransport transport = newFramedTransport(compressed);
+ AutoScalingBufferReadTransport anotherConnection =
+ new AutoScalingBufferReadTransport(100)) {
+ byte[] payload = new byte[1000];
+ new Random(0).nextBytes(payload);
+ transport.write(payload);
+ transport.flush();
+ byte[] writeBuffer = transport.writeBuffer.getBuffer();
+ long retainedMemory = memoryControl.getUsedMemoryInBytes();
+ // Only leave enough quota for the other connection to grow from 100
to 200 bytes.
+ memoryControl.setLimit(retainedMemory + 100);
+ for (int i = 0; i < 3 * (RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1); i++) {
+ setLastShrinkTime(transport.writeBuffer, 0);
+ transport.write(payload, 0, 600);
+ transport.write(payload, 600, 400);
+ transport.flush();
+ Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+ }
+ transport.write(payload, 0, 600);
+ transport.write(payload, 600, 300);
+ Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer());
+ TMemoryBuffer input = new TMemoryBuffer(0);
+ input.write(payload, 0, 200);
+ anotherConnection.fill(input, 200);
+
+ Assert.assertEquals(0,
AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount());
+ Assert.assertEquals(retainedMemory + 100,
memoryControl.getUsedMemoryInBytes());
+ Assert.assertArrayEquals(
+ Arrays.copyOf(payload, 200),
Arrays.copyOf(anotherConnection.getBuffer(), 200));
+ Assert.assertArrayEquals(
+ Arrays.copyOf(payload, 900),
Arrays.copyOf(transport.writeBuffer.getBuffer(), 900));
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ }
+
+ @Test
+ public void testFailedIncrementalGrowthPreservesFrameAndAccounting() throws
Exception {
+ memoryControl.setLimit(120);
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+ AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
+
+ try (AutoScalingBufferWriteTransport transport = new
AutoScalingBufferWriteTransport(100)) {
+ byte[] payload = new byte[130];
+ new Random(0).nextBytes(payload);
+ transport.write(payload, 0, 90);
+ byte[] originalBuffer = transport.getBuffer();
+ try {
+ transport.write(payload, 90, 40);
+ Assert.fail("Expected TTransportException");
+ } catch (TTransportException expected) {
+ Assert.assertSame(originalBuffer, transport.getBuffer());
+ Assert.assertEquals(90, transport.getPos());
+ Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes());
+ Assert.assertArrayEquals(
+ Arrays.copyOf(payload, 90), Arrays.copyOf(transport.getBuffer(),
90));
+ }
+
+ memoryControl.setLimit(130);
+ transport.write(payload, 90, 40);
+ Assert.assertEquals(130, transport.getBuffer().length);
+ Assert.assertEquals(130, transport.getPos());
+ Assert.assertEquals(130, memoryControl.getUsedMemoryInBytes());
+ Assert.assertArrayEquals(payload, transport.getBuffer());
+ Assert.assertEquals(1,
AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount());
+ transport.close();
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+
+ @Test
+ public void testConcurrentReadBuffersEnforceSharedBudget() throws Exception {
+ memoryControl.setLimit(1000);
+ AutoResizingBufferMemoryManager.setMemoryControl(memoryControl);
+ AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
+
+ AutoScalingBufferReadTransport first = new
AutoScalingBufferReadTransport(100);
+ AutoScalingBufferReadTransport second = new
AutoScalingBufferReadTransport(100);
+ AutoScalingBufferReadTransport third = new
AutoScalingBufferReadTransport(100);
+ try {
+ byte[] payload = new byte[400];
+ payload[399] = 42;
+ TMemoryBuffer input = new TMemoryBuffer(0);
+ input.write(payload);
+ first.fill(input, payload.length);
+ input.write(payload);
+ second.fill(input, payload.length);
+ Assert.assertEquals(900, memoryControl.getUsedMemoryInBytes());
+
+ input.write(payload);
+ try {
+ third.fill(input, payload.length);
+ Assert.fail("Expected TTransportException");
+ } catch (TTransportException expected) {
+ Assert.assertEquals(900, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(1,
AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount());
+ }
+
+ first.close();
+ third.fill(input, payload.length);
+ Assert.assertEquals(800, memoryControl.getUsedMemoryInBytes());
+ Assert.assertEquals(42, second.getBuffer()[399]);
+ Assert.assertEquals(42, third.getBuffer()[399]);
+ } finally {
+ first.close();
+ second.close();
+ third.close();
+ }
+ Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes());
+ }
+
@Test
public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception {
memoryControl.setLimit(120);
@@ -67,6 +399,7 @@ public class AutoResizingBufferTest {
AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1);
AutoResizingBuffer buffer = new AutoResizingBuffer(100);
+ buffer.array()[99] = 42;
try {
buffer.resizeIfNecessary(200);
Assert.fail("Expected IOException");
@@ -74,6 +407,8 @@ public class AutoResizingBufferTest {
Assert.assertTrue(e.getMessage().contains("100"));
Assert.assertTrue(e.getMessage().contains("5"));
} finally {
+ Assert.assertEquals(100, buffer.array().length);
+ Assert.assertEquals(42, buffer.array()[99]);
Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes());
buffer.close();
}
@@ -195,6 +530,18 @@ public class AutoResizingBufferTest {
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);
+ }
+
private static class TestMemoryControl implements
AutoResizingBufferMemoryControl {
private final AtomicLong usedMemoryInBytes = new AtomicLong();
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfo.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfo.java
index 5491f113ba7..75c7a73a42b 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfo.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfo.java
@@ -448,8 +448,12 @@ public class SystemInfo {
}
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;
}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfoTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfoTest.java
new file mode 100644
index 00000000000..8c5e3e362d5
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfoTest.java
@@ -0,0 +1,232 @@
+/*
+ * 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.db.storageengine.rescon.memory;
+
+import org.apache.iotdb.db.exception.WriteProcessRejectException;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegionInfo;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable;
+import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class SystemInfoTest {
+
+ private SystemInfo systemInfo;
+ private DataRegionInfo dataRegionInfo;
+ private TsFileProcessor currentProcessor;
+ private TsFileProcessor largerProcessor;
+ private ExecutorService originalFlushSubmitter;
+ private ExecutorService flushSubmitter;
+
+ @Before
+ public void setUp() throws Exception {
+ systemInfo = SystemInfo.getInstance();
+ systemInfo.close();
+ setMemoryThreshold("memorySizeForMemtable", 1000L);
+ setMemoryThreshold("FLUSH_THRESHOLD", 500.0);
+ setMemoryThreshold("REJECT_THRESHOLD", 800.0);
+ Field field =
SystemInfo.class.getDeclaredField("flushTaskSubmitThreadPool");
+ field.setAccessible(true);
+ originalFlushSubmitter = (ExecutorService) field.get(systemInfo);
+ flushSubmitter = mock(ExecutorService.class);
+ field.set(systemInfo, flushSubmitter);
+
+ DataRegion dataRegion = mock(DataRegion.class);
+ when(dataRegion.getDatabaseName()).thenReturn("root.test");
+ dataRegionInfo = new DataRegionInfo(dataRegion);
+ currentProcessor = mock(TsFileProcessor.class);
+ largerProcessor = mock(TsFileProcessor.class);
+ mockWorkMemTable(currentProcessor);
+ mockWorkMemTable(largerProcessor);
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(100L);
+ when(largerProcessor.getWorkMemTableRamCost()).thenReturn(600L);
+ dataRegionInfo.initTsFileProcessorInfo(currentProcessor);
+ dataRegionInfo.initTsFileProcessorInfo(largerProcessor);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ setMemoryThreshold("flushTaskSubmitThreadPool", originalFlushSubmitter);
+ systemInfo.close();
+ systemInfo.loadWriteMemory();
+ }
+
+ @Test
+ public void testRejectedWriterCannotFlushSmallerCurrentProcessor() throws
Exception {
+ dataRegionInfo.addStorageGroupMemCost(900);
+
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ Assert.assertTrue(systemInfo.isRejected());
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ }
+
+ @Test
+ public void testRejectedWriterFirstTabletCannotBypassSelection() throws
Exception {
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L);
+ dataRegionInfo.addStorageGroupMemCost(900);
+
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ }
+
+ @Test
+ public void testRepeatedReportsSubmitEachMemtableOnce() throws Exception {
+ dataRegionInfo.addStorageGroupMemCost(900);
+
+ for (int i = 0; i < 10; i++) {
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ }
+ Assert.assertTrue(systemInfo.isRejected());
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ verify(flushSubmitter).submit(any(Runnable.class));
+ }
+
+ @Test
+ public void testRepeatedReportsForFirstTabletSubmitOnce() throws Exception {
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L);
+ dataRegionInfo.addStorageGroupMemCost(900);
+
+ for (int i = 0; i < 10; i++) {
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ }
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ verify(flushSubmitter).submit(any(Runnable.class));
+ }
+
+ @Test
+ public void testNewMemtableCanBeScheduledAfterPreviousFlush() throws
Exception {
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L);
+ when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L);
+ dataRegionInfo.addStorageGroupMemCost(900);
+ Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+
+ mockWorkMemTable(currentProcessor);
+ Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ verify(currentProcessor, times(2)).setWorkMemTableShouldFlush();
+ verify(largerProcessor, never()).setWorkMemTableShouldFlush();
+ verify(flushSubmitter, times(2)).submit(any(Runnable.class));
+ }
+
+ @Test
+ public void testMarkedCurrentProcessorCannotReachHardMemoryLimit() throws
Exception {
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L);
+ when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L);
+ dataRegionInfo.addStorageGroupMemCost(900);
+ Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ dataRegionInfo.addStorageGroupMemCost(100);
+
+ try {
+ systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor);
+ Assert.fail("Expected WriteProcessRejectException");
+ } catch (WriteProcessRejectException expected) {
+ Assert.assertTrue(systemInfo.isRejected());
+ verify(currentProcessor).setWorkMemTableShouldFlush();
+ verify(flushSubmitter).submit(any(Runnable.class));
+ }
+ }
+
+ @Test
+ public void testNoCurrentMemtableCannotBypassRejection() throws Exception {
+ when(currentProcessor.getWorkMemTable()).thenReturn(null);
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L);
+ dataRegionInfo.addStorageGroupMemCost(900);
+
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ Assert.assertTrue(systemInfo.isRejected());
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ }
+
+ @Test
+ public void testRejectionClearsAfterMemoryIsReleased() throws Exception {
+ dataRegionInfo.addStorageGroupMemCost(900);
+ Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+
+ dataRegionInfo.releaseStorageGroupMemCost(500);
+ systemInfo.resetStorageGroupStatus(dataRegionInfo);
+ Assert.assertFalse(systemInfo.isRejected());
+ Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ verify(flushSubmitter).submit(any(Runnable.class));
+ }
+
+ @Test
+ public void testCurrentProcessorCannotExceedHardMemoryLimit() throws
Exception {
+ when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L);
+ when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L);
+ dataRegionInfo.addStorageGroupMemCost(1100);
+
+ try {
+ systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor);
+ Assert.fail("Expected WriteProcessRejectException");
+ } catch (WriteProcessRejectException expected) {
+ Assert.assertTrue(systemInfo.isRejected());
+ }
+ }
+
+ @Test
+ public void testNormalFlushStillSelectsLargestMemtable() throws Exception {
+ dataRegionInfo.addStorageGroupMemCost(700);
+
+ Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo,
currentProcessor));
+ Assert.assertFalse(systemInfo.isRejected());
+ verify(largerProcessor).setWorkMemTableShouldFlush();
+ verify(currentProcessor, never()).setWorkMemTableShouldFlush();
+ }
+
+ private void setMemoryThreshold(String fieldName, Object value) throws
Exception {
+ Field field = SystemInfo.class.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ field.set(systemInfo, value);
+ }
+
+ private void mockWorkMemTable(TsFileProcessor processor) {
+ IMemTable memTable = mock(IMemTable.class);
+ AtomicBoolean marked = new AtomicBoolean();
+ when(memTable.shouldFlush()).thenAnswer(invocation -> marked.get());
+ when(processor.getWorkMemTable()).thenReturn(memTable);
+ doAnswer(
+ invocation -> {
+ marked.set(true);
+ return null;
+ })
+ .when(processor)
+ .setWorkMemTableShouldFlush();
+ }
+}
