This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new 221f79c5fae [To dev/1.3] Fix RPC buffer resizing and duplicate flush
submissions (#18918) (#18971)
221f79c5fae is described below
commit 221f79c5fae7b63d683783db80657a19b60fb6ea
Author: Caideyipi <[email protected]>
AuthorDate: Thu Oct 22 11:41:24 2026 +0800
[To dev/1.3] Fix RPC buffer resizing and duplicate flush submissions
(#18918) (#18971)
Backport the applicable changes from #18918 to the dev/1.3 buffer and
memory-admission implementations. Incremental serialization only grows
buffers; complete frame sizes drive shrinking in ordinary and compressed
transports. Count marked memtables without resubmitting pending flushes.
Adapt the regression coverage to 1.3 and retain its existing admission
decisions. The RPC memory-budget fallback is omitted because this branch
has no AutoResizingBufferMemoryManager.
(cherry picked from commit 83d419796fb91a6cccb874414820d3865ab34b98)
---
.../org/apache/iotdb/rpc/AutoResizingBuffer.java | 34 +--
.../iotdb/rpc/AutoScalingBufferWriteTransport.java | 4 +-
.../rpc/TCompressedElasticFramedTransport.java | 4 +-
.../apache/iotdb/rpc/TElasticFramedTransport.java | 4 +-
.../apache/iotdb/rpc/AutoResizingBufferTest.java | 221 ++++++++++++++++++++
.../db/storageengine/rescon/memory/SystemInfo.java | 8 +-
.../rescon/memory/SystemInfoTest.java | 232 +++++++++++++++++++++
7 files changed, 484 insertions(+), 23 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 180d05b3d93..ca47d80acae 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
@@ -22,14 +22,15 @@ package org.apache.iotdb.rpc;
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 @@ class AutoResizingBuffer {
}
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();
}
}
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 ae833fad239..9e57e53768f 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
@@ -44,7 +44,9 @@ public class AutoScalingBufferWriteTransport extends
NonOpenTransport {
@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;
}
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 07768dd98bd..ab00ca2ac63 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
@@ -79,9 +79,7 @@ public abstract class TCompressedElasticFramedTransport
extends TElasticFramedTr
}
writeBuffer.reset();
- if (thriftDefaultBufferSize < length) {
- writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
- }
+ writeBuffer.resizeIfNecessary(length);
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 40811a00f41..cbc74e22f42 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
@@ -260,9 +260,7 @@ public class TElasticFramedTransport extends TTransport {
underlying.write(i32buf, 0, 4);
underlying.write(writeBuffer.getBuffer(), 0, length);
writeBuffer.reset();
- if (length > thriftDefaultBufferSize) {
- writeBuffer.resizeIfNecessary(thriftDefaultBufferSize);
- }
+ writeBuffer.resizeIfNecessary(length);
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
new file mode 100644
index 00000000000..d4ddaa00313
--- /dev/null
+++
b/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);
+ }
+}
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 e8cd62b6d16..38d0c75c073 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
@@ -493,8 +493,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..934ea252d5d
--- /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.allocateWriteMemory();
+ }
+
+ @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();
+ }
+}
