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 aa0ea40541ce8..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,14 +23,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 when the memory budget allows it. + * Shrinking: If the required size < current capacity * 0.6, and such small requests last for more + * than 5 times, shrink to the middle of the required size and current capacity. The required size + * is bounded below by the initial capacity. */ class AutoResizingBuffer implements AutoCloseable { @@ -56,23 +57,28 @@ 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); + resize(newCapacity, size); bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; - } else if (size > initialCapacity - && currentCapacity * loadFactor > size - && bufTooLargeCounter-- <= 0 - && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) { - // do not resize if it is reading the request size and do not shrink too often - resize(size + (currentCapacity - size) / 2); - bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; - lastShrinkTime = System.currentTimeMillis(); } } @@ -89,11 +95,14 @@ 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/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 164d0f9e287e2..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 { @@ -60,6 +62,336 @@ 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(); + byte[] largeWriteBuffer = transport.writeBuffer.getBuffer(); + byte[] largeReadBuffer = transport.readBuffer.getBuffer(); + long cooldown = System.currentTimeMillis() + RpcUtils.MIN_SHRINK_INTERVAL; + setLastShrinkTime(transport.writeBuffer, cooldown); + setLastShrinkTime(transport.readBuffer, cooldown); + + for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) { + transport.write(new byte[] {42}); + transport.flush(); + byte[] smallFrame = new byte[1]; + transport.readAll(smallFrame, 0, smallFrame.length); + Assert.assertEquals(42, smallFrame[0]); + } + Assert.assertSame(largeWriteBuffer, transport.writeBuffer.getBuffer()); + Assert.assertSame(largeReadBuffer, transport.readBuffer.getBuffer()); + + setLastShrinkTime(transport.writeBuffer, 0); + setLastShrinkTime(transport.readBuffer, 0); + transport.write(new byte[] {42}); + transport.flush(); + byte[] smallFrame = new byte[1]; + transport.readAll(smallFrame, 0, 1); + Assert.assertEquals(42, smallFrame[0]); + Assert.assertTrue(transport.writeBuffer.getBuffer().length < payload.length); + Assert.assertTrue(transport.readBuffer.getBuffer().length < payload.length); + Assert.assertTrue(transport.writeBuffer.getBuffer().length >= 100); + Assert.assertTrue(transport.readBuffer.getBuffer().length >= 100); + Assert.assertTrue(memoryControl.getUsedMemoryInBytes() < largeFrameMemory); + + byte[] shrunkBuffer = transport.writeBuffer.getBuffer(); + for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) { + transport.write(new byte[] {42}); + transport.flush(); + transport.readAll(smallFrame, 0, 1); + Assert.assertEquals(42, smallFrame[0]); + } + Assert.assertSame(shrunkBuffer, transport.writeBuffer.getBuffer()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + } + + @Test + public void testRepeatedLargeFramesKeepWriteCapacity() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + + for (boolean compressed : new boolean[] {false, true}) { + try (TElasticFramedTransport transport = newFramedTransport(compressed)) { + byte[] payload = new byte[1000]; + new Random(0).nextBytes(payload); + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + byte[] writeBuffer = transport.writeBuffer.getBuffer(); + for (int i = 0; i < 3 * (RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1); i++) { + setLastShrinkTime(transport.writeBuffer, 0); + long allocationCount = AutoResizingBufferMemoryManager.getMemoryAllocationCount(); + transport.write(payload); + transport.flush(); + Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer()); + Assert.assertEquals( + allocationCount, AutoResizingBufferMemoryManager.getMemoryAllocationCount()); + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + } + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + } + + @Test + public void testSmallWritesDoNotShrinkAnUnfinishedFrame() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + + for (boolean compressed : new boolean[] {false, true}) { + try (TElasticFramedTransport transport = newFramedTransport(compressed)) { + byte[] payload = new byte[1000]; + new Random(0).nextBytes(payload); + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + byte[] writeBuffer = transport.writeBuffer.getBuffer(); + setLastShrinkTime(transport.writeBuffer, 0); + + int prefixLength = RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1; + for (int i = 0; i < prefixLength; i++) { + transport.write(payload, i, 1); + Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer()); + Assert.assertEquals(i + 1, transport.writeBuffer.getPos()); + } + transport.write(payload, prefixLength, payload.length - prefixLength); + Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer()); + transport.flush(); + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + } + + @Test + public void testLargeFrameReuseDoesNotStarveAnotherConnection() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + for (boolean compressed : new boolean[] {false, true}) { + memoryControl.setLimit(Long.MAX_VALUE); + try (TElasticFramedTransport transport = newFramedTransport(compressed); + AutoScalingBufferReadTransport anotherConnection = + new AutoScalingBufferReadTransport(100)) { + byte[] payload = new byte[1000]; + new Random(0).nextBytes(payload); + transport.write(payload); + transport.flush(); + byte[] writeBuffer = transport.writeBuffer.getBuffer(); + long retainedMemory = memoryControl.getUsedMemoryInBytes(); + // Only leave enough quota for the other connection to grow from 100 to 200 bytes. + memoryControl.setLimit(retainedMemory + 100); + for (int i = 0; i < 3 * (RpcUtils.MAX_BUFFER_OVERSIZE_TIME + 1); i++) { + setLastShrinkTime(transport.writeBuffer, 0); + transport.write(payload, 0, 600); + transport.write(payload, 600, 400); + transport.flush(); + Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer()); + } + transport.write(payload, 0, 600); + transport.write(payload, 600, 300); + Assert.assertSame(writeBuffer, transport.writeBuffer.getBuffer()); + TMemoryBuffer input = new TMemoryBuffer(0); + input.write(payload, 0, 200); + anotherConnection.fill(input, 200); + + Assert.assertEquals(0, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); + Assert.assertEquals(retainedMemory + 100, memoryControl.getUsedMemoryInBytes()); + Assert.assertArrayEquals( + Arrays.copyOf(payload, 200), Arrays.copyOf(anotherConnection.getBuffer(), 200)); + Assert.assertArrayEquals( + Arrays.copyOf(payload, 900), Arrays.copyOf(transport.writeBuffer.getBuffer(), 900)); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + } + + @Test + public void testFailedIncrementalGrowthPreservesFrameAndAccounting() throws Exception { + memoryControl.setLimit(120); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + try (AutoScalingBufferWriteTransport transport = new AutoScalingBufferWriteTransport(100)) { + byte[] payload = new byte[130]; + new Random(0).nextBytes(payload); + transport.write(payload, 0, 90); + byte[] originalBuffer = transport.getBuffer(); + try { + transport.write(payload, 90, 40); + Assert.fail("Expected TTransportException"); + } catch (TTransportException expected) { + Assert.assertSame(originalBuffer, transport.getBuffer()); + Assert.assertEquals(90, transport.getPos()); + Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes()); + Assert.assertArrayEquals( + Arrays.copyOf(payload, 90), Arrays.copyOf(transport.getBuffer(), 90)); + } + + memoryControl.setLimit(130); + transport.write(payload, 90, 40); + Assert.assertEquals(130, transport.getBuffer().length); + Assert.assertEquals(130, transport.getPos()); + Assert.assertEquals(130, memoryControl.getUsedMemoryInBytes()); + Assert.assertArrayEquals(payload, transport.getBuffer()); + Assert.assertEquals(1, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); + transport.close(); + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testConcurrentReadBuffersEnforceSharedBudget() throws Exception { + memoryControl.setLimit(1000); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + AutoScalingBufferReadTransport first = new AutoScalingBufferReadTransport(100); + AutoScalingBufferReadTransport second = new AutoScalingBufferReadTransport(100); + AutoScalingBufferReadTransport third = new AutoScalingBufferReadTransport(100); + try { + byte[] payload = new byte[400]; + payload[399] = 42; + TMemoryBuffer input = new TMemoryBuffer(0); + input.write(payload); + first.fill(input, payload.length); + input.write(payload); + second.fill(input, payload.length); + Assert.assertEquals(900, memoryControl.getUsedMemoryInBytes()); + + input.write(payload); + try { + third.fill(input, payload.length); + Assert.fail("Expected TTransportException"); + } catch (TTransportException expected) { + Assert.assertEquals(900, memoryControl.getUsedMemoryInBytes()); + Assert.assertEquals(1, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); + } + + first.close(); + third.fill(input, payload.length); + Assert.assertEquals(800, memoryControl.getUsedMemoryInBytes()); + Assert.assertEquals(42, second.getBuffer()[399]); + Assert.assertEquals(42, third.getBuffer()[399]); + } finally { + first.close(); + second.close(); + third.close(); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + @Test public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { memoryControl.setLimit(120); @@ -67,6 +399,7 @@ public 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 +407,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(); } @@ -195,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 5491f113ba7e4..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 @@ -448,8 +448,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..8c5e3e362d56c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/memory/SystemInfoTest.java @@ -0,0 +1,232 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.rescon.memory; + +import org.apache.iotdb.db.exception.WriteProcessRejectException; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.DataRegionInfo; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; +import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class SystemInfoTest { + + private SystemInfo systemInfo; + private DataRegionInfo dataRegionInfo; + private TsFileProcessor currentProcessor; + private TsFileProcessor largerProcessor; + private ExecutorService originalFlushSubmitter; + private ExecutorService flushSubmitter; + + @Before + public void setUp() throws Exception { + systemInfo = SystemInfo.getInstance(); + systemInfo.close(); + setMemoryThreshold("memorySizeForMemtable", 1000L); + setMemoryThreshold("FLUSH_THRESHOLD", 500.0); + setMemoryThreshold("REJECT_THRESHOLD", 800.0); + Field field = SystemInfo.class.getDeclaredField("flushTaskSubmitThreadPool"); + field.setAccessible(true); + originalFlushSubmitter = (ExecutorService) field.get(systemInfo); + flushSubmitter = mock(ExecutorService.class); + field.set(systemInfo, flushSubmitter); + + DataRegion dataRegion = mock(DataRegion.class); + when(dataRegion.getDatabaseName()).thenReturn("root.test"); + dataRegionInfo = new DataRegionInfo(dataRegion); + currentProcessor = mock(TsFileProcessor.class); + largerProcessor = mock(TsFileProcessor.class); + mockWorkMemTable(currentProcessor); + mockWorkMemTable(largerProcessor); + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(100L); + when(largerProcessor.getWorkMemTableRamCost()).thenReturn(600L); + dataRegionInfo.initTsFileProcessorInfo(currentProcessor); + dataRegionInfo.initTsFileProcessorInfo(largerProcessor); + } + + @After + public void tearDown() throws Exception { + setMemoryThreshold("flushTaskSubmitThreadPool", originalFlushSubmitter); + systemInfo.close(); + systemInfo.loadWriteMemory(); + } + + @Test + public void testRejectedWriterCannotFlushSmallerCurrentProcessor() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + } + + @Test + public void testRejectedWriterFirstTabletCannotBypassSelection() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L); + dataRegionInfo.addStorageGroupMemCost(900); + + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + } + + @Test + public void testRepeatedReportsSubmitEachMemtableOnce() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + + for (int i = 0; i < 10; i++) { + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + } + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter).submit(any(Runnable.class)); + } + + @Test + public void testRepeatedReportsForFirstTabletSubmitOnce() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L); + dataRegionInfo.addStorageGroupMemCost(900); + + for (int i = 0; i < 10; i++) { + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + } + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter).submit(any(Runnable.class)); + } + + @Test + public void testNewMemtableCanBeScheduledAfterPreviousFlush() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L); + when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L); + dataRegionInfo.addStorageGroupMemCost(900); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + + mockWorkMemTable(currentProcessor); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(currentProcessor, times(2)).setWorkMemTableShouldFlush(); + verify(largerProcessor, never()).setWorkMemTableShouldFlush(); + verify(flushSubmitter, times(2)).submit(any(Runnable.class)); + } + + @Test + public void testMarkedCurrentProcessorCannotReachHardMemoryLimit() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L); + when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L); + dataRegionInfo.addStorageGroupMemCost(900); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + dataRegionInfo.addStorageGroupMemCost(100); + + try { + systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor); + Assert.fail("Expected WriteProcessRejectException"); + } catch (WriteProcessRejectException expected) { + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor).setWorkMemTableShouldFlush(); + verify(flushSubmitter).submit(any(Runnable.class)); + } + } + + @Test + public void testNoCurrentMemtableCannotBypassRejection() throws Exception { + when(currentProcessor.getWorkMemTable()).thenReturn(null); + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(0L); + dataRegionInfo.addStorageGroupMemCost(900); + + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertTrue(systemInfo.isRejected()); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + verify(largerProcessor).setWorkMemTableShouldFlush(); + } + + @Test + public void testRejectionClearsAfterMemoryIsReleased() throws Exception { + dataRegionInfo.addStorageGroupMemCost(900); + Assert.assertFalse(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + + dataRegionInfo.releaseStorageGroupMemCost(500); + systemInfo.resetStorageGroupStatus(dataRegionInfo); + Assert.assertFalse(systemInfo.isRejected()); + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + verify(flushSubmitter).submit(any(Runnable.class)); + } + + @Test + public void testCurrentProcessorCannotExceedHardMemoryLimit() throws Exception { + when(currentProcessor.getWorkMemTableRamCost()).thenReturn(700L); + when(largerProcessor.getWorkMemTableRamCost()).thenReturn(100L); + dataRegionInfo.addStorageGroupMemCost(1100); + + try { + systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor); + Assert.fail("Expected WriteProcessRejectException"); + } catch (WriteProcessRejectException expected) { + Assert.assertTrue(systemInfo.isRejected()); + } + } + + @Test + public void testNormalFlushStillSelectsLargestMemtable() throws Exception { + dataRegionInfo.addStorageGroupMemCost(700); + + Assert.assertTrue(systemInfo.reportStorageGroupStatus(dataRegionInfo, currentProcessor)); + Assert.assertFalse(systemInfo.isRejected()); + verify(largerProcessor).setWorkMemTableShouldFlush(); + verify(currentProcessor, never()).setWorkMemTableShouldFlush(); + } + + private void setMemoryThreshold(String fieldName, Object value) throws Exception { + Field field = SystemInfo.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(systemInfo, value); + } + + private void mockWorkMemTable(TsFileProcessor processor) { + IMemTable memTable = mock(IMemTable.class); + AtomicBoolean marked = new AtomicBoolean(); + when(memTable.shouldFlush()).thenAnswer(invocation -> marked.get()); + when(processor.getWorkMemTable()).thenReturn(memTable); + doAnswer( + invocation -> { + marked.set(true); + return null; + }) + .when(processor) + .setWorkMemTableShouldFlush(); + } +}