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 180d05b3d9324..ca47d80acae23 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 @@ 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. * *

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(); } } 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 ae833fad23962..9e57e53768f95 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 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; } 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 07768dd98bdd8..ab00ca2ac6393 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 void flush() throws TTransportException { } 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 40811a00f41df..cbc74e22f42c6 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 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(); } 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 0000000000000..d4ddaa003138d --- /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 e8cd62b6d166f..38d0c75c07385 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 @@ 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; } 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 0000000000000..934ea252d5d27 --- /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(); + } +}