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 { @@ -57,20 +58,21 @@ public AutoResizingBuffer(int initialCapacity) throws IOException { public void resizeIfNecessary(int size) throws IOException { reserveCurrentCapacityIfReleased(); + size = Math.max(size, initialCapacity); final int currentCapacity = this.array.length; final double loadFactor = 0.6; 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 + } else if (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); + // Small frames and flush resets must also release oversized capacity, without shrinking + // below the initial capacity or resizing too often. + resize(size + (currentCapacity - size) / 2, size); bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; lastShrinkTime = System.currentTimeMillis(); } @@ -89,11 +91,14 @@ public void close() { } } - private void resize(int newCapacity) throws IOException { + private void resize(int newCapacity, int requiredCapacity) throws IOException { final int currentCapacity = array.length; if (newCapacity > currentCapacity) { - final int delta = newCapacity - currentCapacity; - memoryAccount.allocate(delta); + final int delta = + (int) + memoryAccount.allocate( + newCapacity - currentCapacity, requiredCapacity - currentCapacity); + newCapacity = currentCapacity + delta; try { array = Arrays.copyOf(array, newCapacity); accountedCapacity += delta; 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 d33a15b44cbe4..bbdbc8bda06ed 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 @@ void allocate(long sizeInBytes) throws IOException { AutoResizingBufferMemoryManager.allocate(memoryControl, sizeInBytes); } + long allocate(long preferredSizeInBytes, long requiredSizeInBytes) throws IOException { + // Growth headroom is optional. If it does not fit, only wait for the bytes actually needed + // by the frame instead of rejecting a request that still fits within the memory budget. + if (preferredSizeInBytes > requiredSizeInBytes + && memoryControl.allocate(preferredSizeInBytes)) { + memoryAllocationCount.incrementAndGet(); + return preferredSizeInBytes; + } + allocate(requiredSizeInBytes); + return requiredSizeInBytes; + } + void release(long sizeInBytes) { AutoResizingBufferMemoryManager.release(memoryControl, sizeInBytes); } 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 164d0f9e287e2..988be55a9e30f 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 @@ -60,6 +60,129 @@ public void testAllocateAndReleaseMemoryWhenResizing() throws Exception { 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(); + + 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.assertTrue(transport.writeBuffer.getBuffer().length < payload.length); + Assert.assertTrue(transport.readBuffer.getBuffer().length < payload.length); + Assert.assertTrue(memoryControl.getUsedMemoryInBytes() < largeFrameMemory); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + } + @Test public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { memoryControl.setLimit(120); @@ -67,6 +190,7 @@ public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); AutoResizingBuffer buffer = new AutoResizingBuffer(100); + buffer.array()[99] = 42; try { buffer.resizeIfNecessary(200); Assert.fail("Expected IOException"); @@ -74,6 +198,8 @@ public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { 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(); } 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 5491f113ba7e4..d3ec5a8b40491 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 @@ -137,7 +137,18 @@ public synchronized boolean reportStorageGroupStatus( totalStorageGroupMemCost, REJECT_THRESHOLD); rejected = true; - if (chooseMemTablesToMarkFlush(tsFileProcessor)) { + boolean currentProcessorSelected = chooseMemTablesToMarkFlush(tsFileProcessor); + if (!currentProcessorSelected + && tsFileProcessor != null + && tsFileProcessor.getWorkMemTable() != null) { + // The writer holds its region's write lock. Waiting for another processor in that region + // to flush would prevent the flush submitter from acquiring the same lock. Let this writer + // finish and flush its own memtable, subject to the hard memory limit below. + tsFileProcessor.setWorkMemTableShouldFlush(); + flushTaskSubmitThreadPool.submit(tsFileProcessor::submitAFlushTask); + currentProcessorSelected = true; + } + if (currentProcessorSelected) { if (totalStorageGroupMemCost < memorySizeForMemtable) { return true; } else { 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..4d1da9208ec68 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfoTest.java @@ -0,0 +1,121 @@ +/* + * 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 static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +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; + + @Before + public void setUp() throws Exception { + systemInfo = SystemInfo.getInstance(); + systemInfo.close(); + setMemoryThreshold("memorySizeForMemtable", 1000L); + setMemoryThreshold("FLUSH_THRESHOLD", 500.0); + setMemoryThreshold("REJECT_THRESHOLD", 800.0); + + DataRegion dataRegion = mock(DataRegion.class); + when(dataRegion.getDatabaseName()).thenReturn("root.test"); + dataRegionInfo = new DataRegionInfo(dataRegion); + currentProcessor = mock(TsFileProcessor.class); + largerProcessor = mock(TsFileProcessor.class); + when(currentProcessor.getWorkMemTable()).thenReturn(mock(IMemTable.class)); + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(100L); + when(largerProcessor.getWorkMemTableRamCost()).thenReturn(600L); + dataRegionInfo.initTsFileProcessorInfo(currentProcessor); + dataRegionInfo.initTsFileProcessorInfo(largerProcessor); + } + + @After + public void tearDown() { + systemInfo.close(); + systemInfo.loadWriteMemory(); + } + + @Test + public void testRejectedWriterCanFlushItsCurrentProcessor() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + + // The larger processor needs the region write lock held by the inserting thread. Admission + // must let that thread finish and flush its current memtable instead of waiting for that lock. + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + } + + @Test + public void testRejectedWriterCanFlushItsFirstTablet() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L); + dataRegionInfo.addStorageGroupMemCost(900); + + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(currentProcessor).setWorkMemTableShouldFlush(); + } + + @Test + public void testCurrentProcessorCannotExceedHardMemoryLimit() throws Exception { + 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); + } +} From ab56b846a85090401530c43adaa6872d6ade1dc6 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:48:47 +0800 Subject: [PATCH 2/3] Preserve RPC frame buffers and avoid duplicate flush submissions --- .../apache/iotdb/rpc/AutoResizingBuffer.java | 28 ++- .../rpc/AutoScalingBufferWriteTransport.java | 4 +- .../TCompressedElasticFramedTransport.java | 10 +- .../iotdb/rpc/TElasticFramedTransport.java | 10 +- .../iotdb/rpc/AutoResizingBufferTest.java | 221 ++++++++++++++++++ .../rescon/memory/SystemInfo.java | 14 +- .../rescon/memory/SystemInfoTest.java | 109 ++++++++- 7 files changed, 365 insertions(+), 31 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 61a3a0fefe1e2..d9772985dc614 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,9 +23,9 @@ 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 when the memory budget allows it. @@ -57,24 +57,28 @@ public AutoResizingBuffer(int initialCapacity) throws IOException { } public void resizeIfNecessary(int size) throws IOException { - reserveCurrentCapacityIfReleased(); + ensureCapacity(size); size = Math.max(size, initialCapacity); final int currentCapacity = this.array.length; final double loadFactor = 0.6; + if (currentCapacity * loadFactor > size + && bufTooLargeCounter-- <= 0 + && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) { + resize(size + (currentCapacity - size) / 2, size); + bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; + lastShrinkTime = System.currentTimeMillis(); + } + } + + void ensureCapacity(int size) throws IOException { + reserveCurrentCapacityIfReleased(); + final int currentCapacity = this.array.length; if (currentCapacity < size) { // Increase by a factor of 1.5x int growCapacity = currentCapacity + (currentCapacity >> 1); int newCapacity = Math.max(growCapacity, size); resize(newCapacity, size); bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; - } else if (currentCapacity * loadFactor > size - && bufTooLargeCounter-- <= 0 - && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) { - // Small frames and flush resets must also release oversized capacity, without shrinking - // below the initial capacity or resizing too often. - resize(size + (currentCapacity - size) / 2, size); - 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 24278067fcdff..eab45f3da48f7 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 int read(byte[] buf, int off, int len) { @Override public void write(byte[] toWrite, int off, int len) throws TTransportException { try { - buf.resizeIfNecessary(pos + len); + // A small prefix during serialization does not indicate a small frame. Only the complete + // frame size, supplied after flushing, may trigger shrinking. + buf.ensureCapacity(pos + len); } catch (IOException e) { throw new TTransportException(e); } 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 b4c6c0994f556..36cca36b9b1e5 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 void flush() throws TTransportException { } writeBuffer.reset(); - if (thriftDefaultBufferSize < length) { - try { - writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); - } catch (IOException e) { - throw new TTransportException(e); - } + try { + writeBuffer.resizeIfNecessary(length); + } catch (IOException e) { + throw new TTransportException(e); } underlying.flush(); } 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 339dee595e7b7..e15c9f594248e 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 void flush() throws TTransportException { underlying.write(i32buf, 0, 4); underlying.write(writeBuffer.getBuffer(), 0, length); writeBuffer.reset(); - if (length > thriftDefaultBufferSize) { - try { - writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); - } catch (IOException e) { - throw new TTransportException(e); - } + try { + writeBuffer.resizeIfNecessary(length); + } catch (IOException e) { + throw new TTransportException(e); } underlying.flush(); } 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 988be55a9e30f..18aed44eec3cf 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 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 { @@ -167,6 +169,11 @@ public void testFramedTransportsShrinkAfterLargeFrames() throws Exception { 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}); @@ -175,14 +182,216 @@ public void testFramedTransportsShrinkAfterLargeFrames() throws Exception { 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); @@ -321,6 +530,18 @@ private void setLastShrinkTime(AutoResizingBuffer buffer, long lastShrinkTime) t 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 d3ec5a8b40491..79c44b0db9dcd 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 @@ -144,8 +144,10 @@ public synchronized boolean reportStorageGroupStatus( // The writer holds its region's write lock. Waiting for another processor in that region // to flush would prevent the flush submitter from acquiring the same lock. Let this writer // finish and flush its own memtable, subject to the hard memory limit below. - tsFileProcessor.setWorkMemTableShouldFlush(); - flushTaskSubmitThreadPool.submit(tsFileProcessor::submitAFlushTask); + if (!tsFileProcessor.getWorkMemTable().shouldFlush()) { + tsFileProcessor.setWorkMemTableShouldFlush(); + flushTaskSubmitThreadPool.submit(tsFileProcessor::submitAFlushTask); + } currentProcessorSelected = true; } if (currentProcessorSelected) { @@ -459,8 +461,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 index 4d1da9208ec68..f7edf0748d6df 100644 --- 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 @@ -31,9 +31,14 @@ 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; @@ -43,6 +48,8 @@ public class SystemInfoTest { private DataRegionInfo dataRegionInfo; private TsFileProcessor currentProcessor; private TsFileProcessor largerProcessor; + private ExecutorService originalFlushSubmitter; + private ExecutorService flushSubmitter; @Before public void setUp() throws Exception { @@ -51,13 +58,19 @@ public void setUp() throws Exception { 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); - when(currentProcessor.getWorkMemTable()).thenReturn(mock(IMemTable.class)); + mockWorkMemTable(currentProcessor); + mockWorkMemTable(largerProcessor); when(currentProcessor.getWorkMemTableRamCost()).thenReturn(100L); when(largerProcessor.getWorkMemTableRamCost()).thenReturn(600L); dataRegionInfo.initTsFileProcessorInfo(currentProcessor); @@ -65,7 +78,8 @@ public void setUp() throws Exception { } @After - public void tearDown() { + public void tearDown() throws Exception { + setMemoryThreshold("flushTaskSubmitThreadPool", originalFlushSubmitter); systemInfo.close(); systemInfo.loadWriteMemory(); } @@ -91,6 +105,83 @@ public void testRejectedWriterCanFlushItsFirstTablet() throws Exception { verify(currentProcessor).setWorkMemTableShouldFlush(); } + @Test + public void testRepeatedReportsSubmitEachMemtableOnce() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + + for (int i = 0; i < 10; i++) { + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + } + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter, times(2)).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.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + } + verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + } + + @Test + public void testNewMemtableCanBeScheduledAfterPreviousFlush() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + + mockWorkMemTable(currentProcessor); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(currentProcessor, times(2)).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter, times(3)).submit(any(Runnable.class)); + } + + @Test + public void testMarkedCurrentProcessorCannotReachHardMemoryLimit() throws Exception { + 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, times(2)).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.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + + dataRegionInfo.releaseStorageGroupMemCost(500); + systemInfo.resetStorageGroupStatus(dataRegionInfo); + Assert.assertFalse(systemInfo.isRejected()); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + } + @Test public void testCurrentProcessorCannotExceedHardMemoryLimit() throws Exception { dataRegionInfo.addStorageGroupMemCost(1100); @@ -118,4 +209,18 @@ private void setMemoryThreshold(String fieldName, Object value) throws Exception 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(); + } } From 346d3555ea2a7b4de461ec87ac8dd41db4d56100 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:48:44 +0800 Subject: [PATCH 3/3] Revert fallback flush trigger in memory admission --- .../rescon/memory/SystemInfo.java | 15 +------ .../rescon/memory/SystemInfoTest.java | 44 +++++++++++-------- 2 files changed, 26 insertions(+), 33 deletions(-) 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 79c44b0db9dcd..75c7a73a42b3b 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 @@ -137,20 +137,7 @@ public synchronized boolean reportStorageGroupStatus( totalStorageGroupMemCost, REJECT_THRESHOLD); rejected = true; - boolean currentProcessorSelected = chooseMemTablesToMarkFlush(tsFileProcessor); - if (!currentProcessorSelected - && tsFileProcessor != null - && tsFileProcessor.getWorkMemTable() != null) { - // The writer holds its region's write lock. Waiting for another processor in that region - // to flush would prevent the flush submitter from acquiring the same lock. Let this writer - // finish and flush its own memtable, subject to the hard memory limit below. - if (!tsFileProcessor.getWorkMemTable().shouldFlush()) { - tsFileProcessor.setWorkMemTableShouldFlush(); - flushTaskSubmitThreadPool.submit(tsFileProcessor::submitAFlushTask); - } - currentProcessorSelected = true; - } - if (currentProcessorSelected) { + if (chooseMemTablesToMarkFlush(tsFileProcessor)) { if (totalStorageGroupMemCost < memorySizeForMemtable) { return true; } else { 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 index f7edf0748d6df..8c5e3e362d56c 100644 --- 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 @@ -85,24 +85,23 @@ public void tearDown() throws Exception { } @Test - public void testRejectedWriterCanFlushItsCurrentProcessor() throws Exception { + public void testRejectedWriterCannotFlushSmallerCurrentProcessor() throws Exception { dataRegionInfo.addStorageGroupMemCost(900); - // The larger processor needs the region write lock held by the inserting thread. Admission - // must let that thread finish and flush its current memtable instead of waiting for that lock. - Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); Assert.assertTrue(systemInfo.isRejected()); - verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); verify(largerProcessor).setWorkMemTableShouldFlush(); } @Test - public void testRejectedWriterCanFlushItsFirstTablet() throws Exception { + public void testRejectedWriterFirstTabletCannotBypassSelection() throws Exception { when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L); dataRegionInfo.addStorageGroupMemCost(900); - Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); - verify(currentProcessor).setWorkMemTableShouldFlush(); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); } @Test @@ -110,12 +109,12 @@ public void testRepeatedReportsSubmitEachMemtableOnce() throws Exception { dataRegionInfo.addStorageGroupMemCost(900); for (int i = 0; i < 10; i++) { - Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); } Assert.assertTrue(systemInfo.isRejected()); - verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); verify(largerProcessor).setWorkMemTableShouldFlush(); - verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + verify(flushSubmitter).submit(any(Runnable.class)); } @Test @@ -124,26 +123,31 @@ public void testRepeatedReportsForFirstTabletSubmitOnce() throws Exception { dataRegionInfo.addStorageGroupMemCost(900); for (int i = 0; i < 10; i++) { - Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); } - verify(currentProcessor).setWorkMemTableShouldFlush(); - verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + 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).setWorkMemTableShouldFlush(); - verify(flushSubmitter, times(3)).submit(any(Runnable.class)); + 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); @@ -154,7 +158,7 @@ public void testMarkedCurrentProcessorCannotReachHardMemoryLimit() throws Except } catch (WriteProcessRejectException expected) { Assert.assertTrue(systemInfo.isRejected()); verify(currentProcessor).setWorkMemTableShouldFlush(); - verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + verify(flushSubmitter).submit(any(Runnable.class)); } } @@ -173,17 +177,19 @@ public void testNoCurrentMemtableCannotBypassRejection() throws Exception { @Test public void testRejectionClearsAfterMemoryIsReleased() throws Exception { dataRegionInfo.addStorageGroupMemCost(900); - Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); dataRegionInfo.releaseStorageGroupMemCost(500); systemInfo.resetStorageGroupStatus(dataRegionInfo); Assert.assertFalse(systemInfo.isRejected()); Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); - verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + 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 {