From 3c6b6d901851147024aaf49de62ff4b80dc26490 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:34:51 +0800 Subject: [PATCH 1/2] [Bugfix] Reclaim idle RPC buffers when the shared pool is exhausted --- .../apache/iotdb/rpc/i18n/RpcMessages.java | 4 +- .../apache/iotdb/rpc/i18n/RpcMessages.java | 4 +- .../apache/iotdb/rpc/AutoResizingBuffer.java | 95 ++++-- .../rpc/AutoResizingBufferMemoryControl.java | 10 + .../rpc/AutoResizingBufferMemoryManager.java | 47 ++- .../rpc/AutoScalingBufferReadTransport.java | 10 + .../rpc/AutoScalingBufferWriteTransport.java | 1 + .../TCompressedElasticFramedTransport.java | 3 +- .../iotdb/rpc/TElasticFramedTransport.java | 10 +- .../iotdb/rpc/AutoResizingBufferTest.java | 319 +++++++++++++++++- .../iotdb/commons/memory/MemoryConfig.java | 10 + 11 files changed, 478 insertions(+), 35 deletions(-) diff --git a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java index 9adb0ccf0e0bd..41eadc0698455 100644 --- a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -59,8 +59,8 @@ public final class RpcMessages { // AutoResizingBuffer public static final String AUTO_RESIZING_BUFFER_ALLOCATE_INTERRUPTED = "AutoResizingBuffer was interrupted while allocating %d bytes"; - public static final String AUTO_RESIZING_BUFFER_ALLOCATE_FAILED = - "AutoResizingBuffer failed to allocate %d bytes after %d retries"; + public static final String EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A = + "AutoResizingBuffer failed to allocate %d bytes after %d retries (pool total=%d bytes, used=%d bytes, available=%d bytes)"; // IoTDBRpcDataSet / IoTDBJDBCDataSet public static final String CLOSE_OPERATION_SERVER_ERROR = diff --git a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java index 8cc7d9e1adfd9..5e969e5642158 100644 --- a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -53,8 +53,8 @@ public final class RpcMessages { // AutoResizingBuffer public static final String AUTO_RESIZING_BUFFER_ALLOCATE_INTERRUPTED = "AutoResizingBuffer 分配 %d 字节内存时被中断"; - public static final String AUTO_RESIZING_BUFFER_ALLOCATE_FAILED = - "AutoResizingBuffer 在 %d 次重试后仍无法分配 %d 字节内存"; + public static final String EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A = + "AutoResizingBuffer 无法分配 %d 字节内存，重试 %d 次后仍失败（池总容量=%d 字节，已用=%d 字节，可用=%d 字节）"; // IoTDBRpcDataSet / IoTDBJDBCDataSet public static final String CLOSE_OPERATION_SERVER_ERROR = "服务端关闭操作失败，原因："; 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..a842a00a76ee8 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 @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.Arrays; +import java.util.concurrent.locks.ReentrantLock; /** * Helper class that wraps a byte[] so that it can expand and be reused. Users should call @@ -28,9 +29,10 @@ * needed. * *

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 + * (bounded below by the initial capacity) is less than 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. Unused buffers can also shrink back to their initial capacity under memory pressure. */ class AutoResizingBuffer implements AutoCloseable { @@ -41,6 +43,8 @@ class AutoResizingBuffer implements AutoCloseable { private int accountedCapacity; private boolean closed; private AutoResizingBufferMemoryManager.MemoryAccount memoryAccount; + private final ReentrantLock lock = new ReentrantLock(); + private boolean unused = true; public AutoResizingBuffer(int initialCapacity) throws IOException { this.initialCapacity = initialCapacity; @@ -53,26 +57,66 @@ public AutoResizingBuffer(int initialCapacity) throws IOException { throw e; } this.accountedCapacity = initialCapacity; + AutoResizingBufferMemoryManager.registerBuffer(this); } public void resizeIfNecessary(int size) throws IOException { - reserveCurrentCapacityIfReleased(); - 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); - 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); + lock.lock(); + try { + unused = false; + reserveCurrentCapacityIfReleased(); + final int currentCapacity = this.array.length; + final int requiredCapacity = Math.max(initialCapacity, size); + final double loadFactor = 0.6; + if (currentCapacity < requiredCapacity) { + // Increase by a factor of 1.5x + int growCapacity = currentCapacity + (currentCapacity >> 1); + int newCapacity = Math.max(growCapacity, requiredCapacity); + resize(newCapacity); + bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; + } else if (currentCapacity * loadFactor > requiredCapacity + && bufTooLargeCounter-- <= 0 + && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) { + // Avoid repeatedly copying buffers when frame sizes fluctuate. + resize(requiredCapacity + (currentCapacity - requiredCapacity) / 2); + bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; + lastShrinkTime = System.currentTimeMillis(); + } + } finally { + lock.unlock(); + } + } + + /** The transport must only call this after consuming or sending the entire frame. */ + public void markUnused() { + lock.lock(); + try { + unused = true; + } finally { + lock.unlock(); + } + } + + boolean tryReleaseUnusedMemory(AutoResizingBufferMemoryControl memoryControl) { + // Allocation can already hold another buffer's lock. Never wait for an active buffer. + if (!lock.tryLock()) { + return false; + } + try { + if (closed + || !unused + || array.length <= initialCapacity + || !memoryAccount.usesMemoryControl(memoryControl)) { + return false; + } + array = new byte[initialCapacity]; + memoryAccount.release(accountedCapacity - initialCapacity); + accountedCapacity = initialCapacity; bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; lastShrinkTime = System.currentTimeMillis(); + return true; + } finally { + lock.unlock(); } } @@ -82,10 +126,16 @@ public byte[] array() { @Override public void close() { - if (!closed) { - memoryAccount.release(accountedCapacity); - accountedCapacity = 0; - closed = true; + lock.lock(); + try { + if (!closed) { + memoryAccount.release(accountedCapacity); + accountedCapacity = 0; + closed = true; + AutoResizingBufferMemoryManager.unregisterBuffer(this); + } + } finally { + lock.unlock(); } } @@ -117,6 +167,7 @@ private void reserveCurrentCapacityIfReleased() throws IOException { memoryAccount = newMemoryAccount; accountedCapacity = array.length; closed = false; + AutoResizingBufferMemoryManager.registerBuffer(this); } } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java index 49128727a0b96..fd535e2e681f4 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java @@ -25,4 +25,14 @@ public interface AutoResizingBufferMemoryControl { boolean allocate(long sizeInBytes); void release(long sizeInBytes); + + /** Returns the pool capacity, or -1 if the control does not expose it. */ + default long getTotalMemorySizeInBytes() { + return -1; + } + + /** Returns the pool usage, or -1 if the control does not expose it. */ + default long getUsedMemoryInBytes() { + return -1; + } } 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..93ccd29582736 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 @@ -22,7 +22,10 @@ import org.apache.iotdb.rpc.i18n.RpcMessages; import java.io.IOException; +import java.util.Collections; import java.util.Objects; +import java.util.Set; +import java.util.WeakHashMap; import java.util.concurrent.atomic.AtomicLong; public final class AutoResizingBufferMemoryManager { @@ -47,6 +50,9 @@ public void release(long sizeInBytes) { DEFAULT_MEMORY_ALLOCATE_RETRY_INTERVAL_IN_MS; private static final AtomicLong memoryAllocationCount = new AtomicLong(); private static final AtomicLong memoryAllocationFailureCount = new AtomicLong(); + // Do not keep transports alive just so their unused buffers can be reclaimed. + private static final Set buffers = + Collections.newSetFromMap(new WeakHashMap<>()); private AutoResizingBufferMemoryManager() { // Utility class. @@ -72,13 +78,38 @@ static MemoryAccount createMemoryAccount() { return new MemoryAccount(memoryControl); } + static void registerBuffer(AutoResizingBuffer buffer) { + synchronized (buffers) { + buffers.add(buffer); + } + } + + static void unregisterBuffer(AutoResizingBuffer buffer) { + synchronized (buffers) { + buffers.remove(buffer); + } + } + + private static boolean releaseUnusedBuffers(AutoResizingBufferMemoryControl memoryControl) { + final AutoResizingBuffer[] candidates; + synchronized (buffers) { + candidates = buffers.toArray(new AutoResizingBuffer[0]); + } + boolean released = false; + for (AutoResizingBuffer buffer : candidates) { + released |= buffer.tryReleaseUnusedMemory(memoryControl); + } + return released; + } + private static void allocate(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes) throws IOException { if (sizeInBytes <= 0) { return; } for (int i = 0; i < MEMORY_ALLOCATE_MAX_RETRIES; i++) { - if (memoryControl.allocate(sizeInBytes)) { + if (memoryControl.allocate(sizeInBytes) + || (releaseUnusedBuffers(memoryControl) && memoryControl.allocate(sizeInBytes))) { memoryAllocationCount.incrementAndGet(); return; } @@ -91,11 +122,17 @@ private static void allocate(AutoResizingBufferMemoryControl memoryControl, long } } memoryAllocationFailureCount.incrementAndGet(); + final long totalMemory = memoryControl.getTotalMemorySizeInBytes(); + final long usedMemory = memoryControl.getUsedMemoryInBytes(); throw new IOException( String.format( - RpcMessages.AUTO_RESIZING_BUFFER_ALLOCATE_FAILED, + RpcMessages + .EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A, sizeInBytes, - MEMORY_ALLOCATE_MAX_RETRIES)); + MEMORY_ALLOCATE_MAX_RETRIES, + totalMemory, + usedMemory, + totalMemory < 0 || usedMemory < 0 ? -1 : totalMemory - usedMemory)); } private static void release(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes) { @@ -127,5 +164,9 @@ void allocate(long sizeInBytes) throws IOException { void release(long sizeInBytes) { AutoResizingBufferMemoryManager.release(memoryControl, sizeInBytes); } + + boolean usesMemoryControl(AutoResizingBufferMemoryControl memoryControl) { + return this.memoryControl == memoryControl; + } } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java index 035f8753de1ab..b0612133dab7d 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java @@ -44,11 +44,18 @@ public void fill(TTransport inTrans, int length) throws TTransportException { inTrans.readAll(buf.array(), 0, length); pos = 0; limit = length; + if (length == 0) { + buf.markUnused(); + } } @Override public final int read(byte[] target, int off, int len) { int amtToRead = Math.min(len, getBytesRemainingInBuffer()); + if (amtToRead == 0) { + // A fully consumed buffer may have been reclaimed since the previous frame. + return 0; + } System.arraycopy(buf.array(), pos, target, off, amtToRead); consumeBuffer(amtToRead); return amtToRead; @@ -62,6 +69,9 @@ public void write(byte[] buf, int off, int len) { @Override public final void consumeBuffer(int len) { pos += len; + if (pos == limit) { + buf.markUnused(); + } } @Override 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..d6ffe85a34d04 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 @@ -61,6 +61,7 @@ public int getPos() { public void reset() { pos = 0; + buf.markUnused(); } @Override 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..b441a899aad77 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 @@ -99,7 +99,6 @@ public void flush() throws TTransportException { throw new TTransportException(e); } - writeBuffer.reset(); if (thriftDefaultBufferSize < length) { try { writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); @@ -107,6 +106,8 @@ public void flush() throws TTransportException { throw new TTransportException(e); } } + writeBuffer.reset(); + writeCompressBuffer.reset(); 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..33ea5599410df 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 @@ -139,7 +139,13 @@ public boolean isOpen() { @Override public void open() throws TTransportException { - underlying.open(); + try { + underlying.open(); + } catch (TTransportException e) { + // A failed client constructor may never expose this transport to its caller for cleanup. + close(); + throw e; + } } @Override @@ -316,7 +322,6 @@ public void flush() throws TTransportException { TFramedTransport.encodeFrameSize(length, i32buf); underlying.write(i32buf, 0, 4); underlying.write(writeBuffer.getBuffer(), 0, length); - writeBuffer.reset(); if (length > thriftDefaultBufferSize) { try { writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); @@ -324,6 +329,7 @@ public void flush() throws TTransportException { throw new TTransportException(e); } } + writeBuffer.reset(); 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..c3f9fa0619d25 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 @@ -19,14 +19,33 @@ package org.apache.iotdb.rpc; +import org.apache.iotdb.rpc.i18n.RpcMessages; + +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.server.TServer; +import org.apache.thrift.server.TThreadPoolServer; import org.apache.thrift.transport.TMemoryBuffer; +import org.apache.thrift.transport.TServerSocket; +import org.apache.thrift.transport.TSocket; import org.apache.thrift.transport.TTransportException; +import org.apache.thrift.transport.layered.TFramedTransport; import org.junit.After; import org.junit.Assert; import org.junit.Test; import java.io.IOException; import java.lang.reflect.Field; +import java.net.InetSocketAddress; +import java.net.ServerSocket; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; public class AutoResizingBufferTest { @@ -71,8 +90,16 @@ public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { buffer.resizeIfNecessary(200); Assert.fail("Expected IOException"); } catch (IOException e) { - Assert.assertTrue(e.getMessage().contains("100")); - Assert.assertTrue(e.getMessage().contains("5")); + Assert.assertEquals( + String.format( + RpcMessages + .EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A, + 100, + 5, + 120, + 100, + 20), + e.getMessage()); } finally { Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes()); buffer.close(); @@ -100,6 +127,286 @@ public void testMemoryAllocationMetrics() throws Exception { } } + @Test + public void testSmallFramesShrinkBufferWithoutGoingBelowInitialCapacity() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + + try (AutoResizingBuffer buffer = new AutoResizingBuffer(100)) { + buffer.resizeIfNecessary(1600); + for (int size : new int[] {100, 4, 0}) { + int previousCapacity = buffer.array().length; + setLastShrinkTime(buffer, System.currentTimeMillis() - RpcUtils.MIN_SHRINK_INTERVAL - 1); + for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) { + buffer.resizeIfNecessary(size); + } + Assert.assertTrue(buffer.array().length < previousCapacity); + Assert.assertTrue(buffer.array().length >= 100); + Assert.assertEquals(buffer.array().length, memoryControl.getUsedMemoryInBytes()); + } + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testNewConnectionReclaimsIdleBuffersUnderMemoryPressure() throws Exception { + memoryControl.setLimit(80); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + byte[] payload = new byte[40]; + Arrays.fill(payload, (byte) 7); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { + transport.write(payload); + transport.flush(); + byte[] received = new byte[payload.length]; + transport.readAll(received, 0, received.length); + Assert.assertArrayEquals(payload, received); + Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); + + try (TElasticFramedTransport newConnection = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { + Assert.assertTrue(memoryControl.getUsedMemoryInBytes() <= 80); + Assert.assertTrue( + transport.readBuffer.getBuffer().length == 8 + || transport.writeBuffer.getBuffer().length == 8); + Assert.assertEquals(0, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); + transport.write(payload); + transport.flush(); + transport.readAll(received, 0, received.length); + Assert.assertArrayEquals(payload, received); + } + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test(timeout = 20000) + public void testThriftServerAcceptsNewConnectionAfterLargeFrame() throws Exception { + ServerSocket socket = new ServerSocket(); + socket.bind(new InetSocketAddress("127.0.0.1", 0)); + ExecutorService workers = Executors.newFixedThreadPool(2); + TServer server = + new TThreadPoolServer( + new TThreadPoolServer.Args(new TServerSocket(socket)) + .executorService(workers) + .transportFactory(new TElasticFramedTransport.Factory(8, 1024, true)) + .protocolFactory(new TBinaryProtocol.Factory()) + .processor( + (input, output) -> { + output.writeBinary(input.readBinary()); + output.getTransport().flush(); + })); + Thread serverThread = new Thread(server::serve, "rpc-buffer-memory-test-server"); + serverThread.setDaemon(true); + + // Client buffers are outside the server's memory control, as in a real deployment. + try (TElasticFramedTransport first = + new TElasticFramedTransport( + new TSocket("127.0.0.1", socket.getLocalPort(), 5000), 8, 1024, true); + TElasticFramedTransport second = + new TElasticFramedTransport( + new TSocket("127.0.0.1", socket.getLocalPort(), 5000), 8, 1024, true)) { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + serverThread.start(); + first.open(); + TProtocol firstProtocol = new TBinaryProtocol(first); + byte[] payload = new byte[40]; + Arrays.fill(payload, (byte) 13); + firstProtocol.writeBinary(ByteBuffer.wrap(payload)); + first.flush(); + Assert.assertEquals(ByteBuffer.wrap(payload), firstProtocol.readBinary()); + + memoryControl.setLimit(memoryControl.getUsedMemoryInBytes()); + second.open(); + TProtocol secondProtocol = new TBinaryProtocol(second); + secondProtocol.writeBinary(ByteBuffer.wrap(new byte[] {1, 2, 3, 4})); + second.flush(); + Assert.assertEquals(ByteBuffer.wrap(new byte[] {1, 2, 3, 4}), secondProtocol.readBinary()); + Assert.assertEquals(0, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); + } finally { + server.stop(); + workers.shutdownNow(); + Assert.assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + serverThread.join(5000); + Assert.assertFalse(serverThread.isAlive()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testMemoryPressureDoesNotReclaimUnflushedWriteBuffer() throws Exception { + memoryControl.setLimit(48); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + byte[] payload = new byte[40]; + Arrays.fill(payload, (byte) 9); + TMemoryBuffer wire = new TMemoryBuffer(0); + try (TElasticFramedTransport transport = new TElasticFramedTransport(wire, 8, 1024, true)) { + transport.write(payload); + Assert.assertThrows(IOException.class, () -> new AutoResizingBuffer(8)); + Assert.assertEquals(48, memoryControl.getUsedMemoryInBytes()); + + transport.flush(); + byte[] frame = new byte[4 + payload.length]; + wire.readAll(frame, 0, frame.length); + Assert.assertEquals(payload.length, TFramedTransport.decodeFrameSize(frame)); + Assert.assertArrayEquals(payload, Arrays.copyOfRange(frame, 4, frame.length)); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testMemoryPressurePreservesUnreadFrame() throws Exception { + memoryControl.setLimit(80); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + byte[] payload = new byte[40]; + Arrays.fill(payload, (byte) 11); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, false)) { + transport.write(payload); + transport.flush(); + byte[] received = new byte[payload.length]; + transport.readAll(received, 0, 1); + Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); + + try (TElasticFramedTransport newConnection = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { + Assert.assertEquals(40, transport.readBuffer.getBuffer().length); + Assert.assertEquals(8, transport.writeBuffer.getBuffer().length); + transport.readAll(received, 1, received.length - 1); + Assert.assertArrayEquals(payload, received); + } + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testSnappyConnectionReclaimsIdleCompressionBuffers() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + byte[] payload = new byte[4096]; + new Random(0).nextBytes(payload); + try (TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 32, 65536, true)) { + transport.write(payload); + transport.flush(); + byte[] received = new byte[payload.length]; + transport.readAll(received, 0, received.length); + Assert.assertArrayEquals(payload, received); + long retainedMemory = memoryControl.getUsedMemoryInBytes(); + memoryControl.setLimit(retainedMemory); + + try (TSnappyElasticFramedTransport newConnection = + new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 32, 65536, true)) { + Assert.assertTrue(memoryControl.getUsedMemoryInBytes() < retainedMemory); + memoryControl.setLimit(Long.MAX_VALUE); + transport.write(payload); + transport.flush(); + transport.readAll(received, 0, received.length); + Assert.assertArrayEquals(payload, received); + } + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testConcurrentAllocationsDoNotWaitForEachOthersBufferLocks() throws Exception { + memoryControl.setLimit(80); + CountDownLatch rejectedAllocations = new CountDownLatch(2); + AutoResizingBufferMemoryManager.setMemoryControl( + new AutoResizingBufferMemoryControl() { + @Override + public boolean allocate(long sizeInBytes) { + if (memoryControl.allocate(sizeInBytes)) { + return true; + } + // Both threads now hold their buffer locks while attempting to reclaim memory. + rejectedAllocations.countDown(); + try { + Assert.assertTrue(rejectedAllocations.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + return false; + } + + @Override + public void release(long sizeInBytes) { + memoryControl.release(sizeInBytes); + } + }); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + ExecutorService executor = Executors.newFixedThreadPool(2); + try (AutoResizingBuffer first = new AutoResizingBuffer(8); + AutoResizingBuffer second = new AutoResizingBuffer(8)) { + first.resizeIfNecessary(40); + second.resizeIfNecessary(40); + Future firstFailure = + executor.submit( + () -> Assert.assertThrows(IOException.class, () -> first.resizeIfNecessary(41))); + Future secondFailure = + executor.submit( + () -> Assert.assertThrows(IOException.class, () -> second.resizeIfNecessary(41))); + Assert.assertNotNull(firstFailure.get(10, TimeUnit.SECONDS)); + Assert.assertNotNull(secondFailure.get(10, TimeUnit.SECONDS)); + Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); + } finally { + executor.shutdownNow(); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testMemoryPressureDoesNotReclaimBuffersFromAnotherControl() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + try (AutoScalingBufferWriteTransport original = new AutoScalingBufferWriteTransport(8)) { + original.write(new byte[40]); + original.reset(); + + TestMemoryControl replacement = new TestMemoryControl(); + replacement.setLimit(0); + AutoResizingBufferMemoryManager.setMemoryControl(replacement); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + Assert.assertThrows(IOException.class, () -> new AutoResizingBuffer(8)); + Assert.assertEquals(40, original.getBuffer().length); + Assert.assertEquals(40, memoryControl.getUsedMemoryInBytes()); + Assert.assertEquals(0, replacement.getUsedMemoryInBytes()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testFailedOpenDoesNotLeakBufferMemoryWhenReconnecting() throws Exception { + memoryControl.setLimit(200); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + + for (int i = 0; i < 10; i++) { + TSocket socket = + new TSocket("unavailable", 1) { + @Override + public void open() throws TTransportException { + throw new TTransportException("Simulated connection failure"); + } + }; + TElasticFramedTransport transport = new TElasticFramedTransport(socket, 50, 1024, true); + Assert.assertThrows(TTransportException.class, transport::open); + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + try (TElasticFramedTransport recovered = + new TElasticFramedTransport(new TMemoryBuffer(0), 50, 1024, true)) { + recovered.open(); + Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + @Test public void testElasticFramedTransportReleasesMemoryWhenClosed() throws Exception { AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); @@ -219,10 +526,16 @@ public void release(long sizeInBytes) { usedMemoryInBytes.addAndGet(-sizeInBytes); } - private long getUsedMemoryInBytes() { + @Override + public long getUsedMemoryInBytes() { return usedMemoryInBytes.get(); } + @Override + public long getTotalMemorySizeInBytes() { + return limit; + } + private void setLimit(long limit) { this.limit = limit; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java index 52b08556e3f02..76fc2ba65b5ff 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java @@ -64,6 +64,16 @@ public boolean allocate(long sizeInBytes) { public void release(long sizeInBytes) { memoryBlock.release(sizeInBytes); } + + @Override + public long getTotalMemorySizeInBytes() { + return memoryBlock.getTotalMemorySizeInBytes(); + } + + @Override + public long getUsedMemoryInBytes() { + return memoryBlock.getUsedMemoryInBytes(); + } }); } From d042e49aef2d3a9bc5274cc98be0ebcb849e66ed Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:31:00 +0800 Subject: [PATCH 2/2] [Bugfix] Evict idle RPC clients to release buffers --- .../apache/iotdb/rpc/i18n/RpcMessages.java | 4 +- .../apache/iotdb/rpc/i18n/RpcMessages.java | 4 +- .../apache/iotdb/rpc/AutoResizingBuffer.java | 95 ++--- .../rpc/AutoResizingBufferMemoryControl.java | 10 - .../rpc/AutoResizingBufferMemoryManager.java | 47 +-- .../rpc/AutoScalingBufferReadTransport.java | 10 - .../rpc/AutoScalingBufferWriteTransport.java | 1 - .../TCompressedElasticFramedTransport.java | 3 +- .../iotdb/rpc/TElasticFramedTransport.java | 10 +- .../iotdb/rpc/AutoResizingBufferTest.java | 319 +---------------- .../client/property/ClientPoolProperty.java | 5 +- .../iotdb/commons/memory/MemoryConfig.java | 10 - .../iotdb/rpc/ClientPoolIdleEvictionTest.java | 326 ++++++++++++++++++ 13 files changed, 365 insertions(+), 479 deletions(-) create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/rpc/ClientPoolIdleEvictionTest.java diff --git a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java index 41eadc0698455..9adb0ccf0e0bd 100644 --- a/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/en/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -59,8 +59,8 @@ public final class RpcMessages { // AutoResizingBuffer public static final String AUTO_RESIZING_BUFFER_ALLOCATE_INTERRUPTED = "AutoResizingBuffer was interrupted while allocating %d bytes"; - public static final String EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A = - "AutoResizingBuffer failed to allocate %d bytes after %d retries (pool total=%d bytes, used=%d bytes, available=%d bytes)"; + public static final String AUTO_RESIZING_BUFFER_ALLOCATE_FAILED = + "AutoResizingBuffer failed to allocate %d bytes after %d retries"; // IoTDBRpcDataSet / IoTDBJDBCDataSet public static final String CLOSE_OPERATION_SERVER_ERROR = diff --git a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java index 5e969e5642158..8cc7d9e1adfd9 100644 --- a/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java +++ b/iotdb-client/service-rpc/src/main/i18n/zh/org/apache/iotdb/rpc/i18n/RpcMessages.java @@ -53,8 +53,8 @@ public final class RpcMessages { // AutoResizingBuffer public static final String AUTO_RESIZING_BUFFER_ALLOCATE_INTERRUPTED = "AutoResizingBuffer 分配 %d 字节内存时被中断"; - public static final String EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A = - "AutoResizingBuffer 无法分配 %d 字节内存，重试 %d 次后仍失败（池总容量=%d 字节，已用=%d 字节，可用=%d 字节）"; + public static final String AUTO_RESIZING_BUFFER_ALLOCATE_FAILED = + "AutoResizingBuffer 在 %d 次重试后仍无法分配 %d 字节内存"; // IoTDBRpcDataSet / IoTDBJDBCDataSet public static final String CLOSE_OPERATION_SERVER_ERROR = "服务端关闭操作失败，原因："; 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 a842a00a76ee8..aa0ea40541ce8 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 @@ -21,7 +21,6 @@ import java.io.IOException; import java.util.Arrays; -import java.util.concurrent.locks.ReentrantLock; /** * Helper class that wraps a byte[] so that it can expand and be reused. Users should call @@ -29,10 +28,9 @@ * needed. * *

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 the required size - * (bounded below by the initial capacity) is less than 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. Unused buffers can also shrink back to their initial capacity under memory pressure. + * 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. */ class AutoResizingBuffer implements AutoCloseable { @@ -43,8 +41,6 @@ class AutoResizingBuffer implements AutoCloseable { private int accountedCapacity; private boolean closed; private AutoResizingBufferMemoryManager.MemoryAccount memoryAccount; - private final ReentrantLock lock = new ReentrantLock(); - private boolean unused = true; public AutoResizingBuffer(int initialCapacity) throws IOException { this.initialCapacity = initialCapacity; @@ -57,66 +53,26 @@ public AutoResizingBuffer(int initialCapacity) throws IOException { throw e; } this.accountedCapacity = initialCapacity; - AutoResizingBufferMemoryManager.registerBuffer(this); } public void resizeIfNecessary(int size) throws IOException { - lock.lock(); - try { - unused = false; - reserveCurrentCapacityIfReleased(); - final int currentCapacity = this.array.length; - final int requiredCapacity = Math.max(initialCapacity, size); - final double loadFactor = 0.6; - if (currentCapacity < requiredCapacity) { - // Increase by a factor of 1.5x - int growCapacity = currentCapacity + (currentCapacity >> 1); - int newCapacity = Math.max(growCapacity, requiredCapacity); - resize(newCapacity); - bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; - } else if (currentCapacity * loadFactor > requiredCapacity - && bufTooLargeCounter-- <= 0 - && System.currentTimeMillis() - lastShrinkTime > RpcUtils.MIN_SHRINK_INTERVAL) { - // Avoid repeatedly copying buffers when frame sizes fluctuate. - resize(requiredCapacity + (currentCapacity - requiredCapacity) / 2); - bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; - lastShrinkTime = System.currentTimeMillis(); - } - } finally { - lock.unlock(); - } - } - - /** The transport must only call this after consuming or sending the entire frame. */ - public void markUnused() { - lock.lock(); - try { - unused = true; - } finally { - lock.unlock(); - } - } - - boolean tryReleaseUnusedMemory(AutoResizingBufferMemoryControl memoryControl) { - // Allocation can already hold another buffer's lock. Never wait for an active buffer. - if (!lock.tryLock()) { - return false; - } - try { - if (closed - || !unused - || array.length <= initialCapacity - || !memoryAccount.usesMemoryControl(memoryControl)) { - return false; - } - array = new byte[initialCapacity]; - memoryAccount.release(accountedCapacity - initialCapacity); - accountedCapacity = initialCapacity; + reserveCurrentCapacityIfReleased(); + 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); + 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(); - return true; - } finally { - lock.unlock(); } } @@ -126,16 +82,10 @@ public byte[] array() { @Override public void close() { - lock.lock(); - try { - if (!closed) { - memoryAccount.release(accountedCapacity); - accountedCapacity = 0; - closed = true; - AutoResizingBufferMemoryManager.unregisterBuffer(this); - } - } finally { - lock.unlock(); + if (!closed) { + memoryAccount.release(accountedCapacity); + accountedCapacity = 0; + closed = true; } } @@ -167,7 +117,6 @@ private void reserveCurrentCapacityIfReleased() throws IOException { memoryAccount = newMemoryAccount; accountedCapacity = array.length; closed = false; - AutoResizingBufferMemoryManager.registerBuffer(this); } } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java index fd535e2e681f4..49128727a0b96 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBufferMemoryControl.java @@ -25,14 +25,4 @@ public interface AutoResizingBufferMemoryControl { boolean allocate(long sizeInBytes); void release(long sizeInBytes); - - /** Returns the pool capacity, or -1 if the control does not expose it. */ - default long getTotalMemorySizeInBytes() { - return -1; - } - - /** Returns the pool usage, or -1 if the control does not expose it. */ - default long getUsedMemoryInBytes() { - return -1; - } } 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 93ccd29582736..d33a15b44cbe4 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 @@ -22,10 +22,7 @@ import org.apache.iotdb.rpc.i18n.RpcMessages; import java.io.IOException; -import java.util.Collections; import java.util.Objects; -import java.util.Set; -import java.util.WeakHashMap; import java.util.concurrent.atomic.AtomicLong; public final class AutoResizingBufferMemoryManager { @@ -50,9 +47,6 @@ public void release(long sizeInBytes) { DEFAULT_MEMORY_ALLOCATE_RETRY_INTERVAL_IN_MS; private static final AtomicLong memoryAllocationCount = new AtomicLong(); private static final AtomicLong memoryAllocationFailureCount = new AtomicLong(); - // Do not keep transports alive just so their unused buffers can be reclaimed. - private static final Set buffers = - Collections.newSetFromMap(new WeakHashMap<>()); private AutoResizingBufferMemoryManager() { // Utility class. @@ -78,38 +72,13 @@ static MemoryAccount createMemoryAccount() { return new MemoryAccount(memoryControl); } - static void registerBuffer(AutoResizingBuffer buffer) { - synchronized (buffers) { - buffers.add(buffer); - } - } - - static void unregisterBuffer(AutoResizingBuffer buffer) { - synchronized (buffers) { - buffers.remove(buffer); - } - } - - private static boolean releaseUnusedBuffers(AutoResizingBufferMemoryControl memoryControl) { - final AutoResizingBuffer[] candidates; - synchronized (buffers) { - candidates = buffers.toArray(new AutoResizingBuffer[0]); - } - boolean released = false; - for (AutoResizingBuffer buffer : candidates) { - released |= buffer.tryReleaseUnusedMemory(memoryControl); - } - return released; - } - private static void allocate(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes) throws IOException { if (sizeInBytes <= 0) { return; } for (int i = 0; i < MEMORY_ALLOCATE_MAX_RETRIES; i++) { - if (memoryControl.allocate(sizeInBytes) - || (releaseUnusedBuffers(memoryControl) && memoryControl.allocate(sizeInBytes))) { + if (memoryControl.allocate(sizeInBytes)) { memoryAllocationCount.incrementAndGet(); return; } @@ -122,17 +91,11 @@ private static void allocate(AutoResizingBufferMemoryControl memoryControl, long } } memoryAllocationFailureCount.incrementAndGet(); - final long totalMemory = memoryControl.getTotalMemorySizeInBytes(); - final long usedMemory = memoryControl.getUsedMemoryInBytes(); throw new IOException( String.format( - RpcMessages - .EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A, + RpcMessages.AUTO_RESIZING_BUFFER_ALLOCATE_FAILED, sizeInBytes, - MEMORY_ALLOCATE_MAX_RETRIES, - totalMemory, - usedMemory, - totalMemory < 0 || usedMemory < 0 ? -1 : totalMemory - usedMemory)); + MEMORY_ALLOCATE_MAX_RETRIES)); } private static void release(AutoResizingBufferMemoryControl memoryControl, long sizeInBytes) { @@ -164,9 +127,5 @@ void allocate(long sizeInBytes) throws IOException { void release(long sizeInBytes) { AutoResizingBufferMemoryManager.release(memoryControl, sizeInBytes); } - - boolean usesMemoryControl(AutoResizingBufferMemoryControl memoryControl) { - return this.memoryControl == memoryControl; - } } } diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java index b0612133dab7d..035f8753de1ab 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferReadTransport.java @@ -44,18 +44,11 @@ public void fill(TTransport inTrans, int length) throws TTransportException { inTrans.readAll(buf.array(), 0, length); pos = 0; limit = length; - if (length == 0) { - buf.markUnused(); - } } @Override public final int read(byte[] target, int off, int len) { int amtToRead = Math.min(len, getBytesRemainingInBuffer()); - if (amtToRead == 0) { - // A fully consumed buffer may have been reclaimed since the previous frame. - return 0; - } System.arraycopy(buf.array(), pos, target, off, amtToRead); consumeBuffer(amtToRead); return amtToRead; @@ -69,9 +62,6 @@ public void write(byte[] buf, int off, int len) { @Override public final void consumeBuffer(int len) { pos += len; - if (pos == limit) { - buf.markUnused(); - } } @Override 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 d6ffe85a34d04..24278067fcdff 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 @@ -61,7 +61,6 @@ public int getPos() { public void reset() { pos = 0; - buf.markUnused(); } @Override 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 b441a899aad77..b4c6c0994f556 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 @@ -99,6 +99,7 @@ public void flush() throws TTransportException { throw new TTransportException(e); } + writeBuffer.reset(); if (thriftDefaultBufferSize < length) { try { writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); @@ -106,8 +107,6 @@ public void flush() throws TTransportException { throw new TTransportException(e); } } - writeBuffer.reset(); - writeCompressBuffer.reset(); 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 33ea5599410df..339dee595e7b7 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 @@ -139,13 +139,7 @@ public boolean isOpen() { @Override public void open() throws TTransportException { - try { - underlying.open(); - } catch (TTransportException e) { - // A failed client constructor may never expose this transport to its caller for cleanup. - close(); - throw e; - } + underlying.open(); } @Override @@ -322,6 +316,7 @@ public void flush() throws TTransportException { TFramedTransport.encodeFrameSize(length, i32buf); underlying.write(i32buf, 0, 4); underlying.write(writeBuffer.getBuffer(), 0, length); + writeBuffer.reset(); if (length > thriftDefaultBufferSize) { try { writeBuffer.resizeIfNecessary(thriftDefaultBufferSize); @@ -329,7 +324,6 @@ public void flush() throws TTransportException { throw new TTransportException(e); } } - writeBuffer.reset(); 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 c3f9fa0619d25..164d0f9e287e2 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 @@ -19,33 +19,14 @@ package org.apache.iotdb.rpc; -import org.apache.iotdb.rpc.i18n.RpcMessages; - -import org.apache.thrift.protocol.TBinaryProtocol; -import org.apache.thrift.protocol.TProtocol; -import org.apache.thrift.server.TServer; -import org.apache.thrift.server.TThreadPoolServer; import org.apache.thrift.transport.TMemoryBuffer; -import org.apache.thrift.transport.TServerSocket; -import org.apache.thrift.transport.TSocket; import org.apache.thrift.transport.TTransportException; -import org.apache.thrift.transport.layered.TFramedTransport; import org.junit.After; import org.junit.Assert; import org.junit.Test; import java.io.IOException; import java.lang.reflect.Field; -import java.net.InetSocketAddress; -import java.net.ServerSocket; -import java.nio.ByteBuffer; -import java.util.Arrays; -import java.util.Random; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; public class AutoResizingBufferTest { @@ -90,16 +71,8 @@ public void testThrowIOExceptionWhenMemoryIsInsufficient() throws Exception { buffer.resizeIfNecessary(200); Assert.fail("Expected IOException"); } catch (IOException e) { - Assert.assertEquals( - String.format( - RpcMessages - .EXCEPTION_AUTORESIZINGBUFFER_FAILED_TO_ALLOCATE_ARG_BYTES_AFTER_ARG_RETRIES_POOL_TOTAL_ARG_BYTES_USED_ARG_BYTES_AVAILABLE_ARG_BYTES_8244057A, - 100, - 5, - 120, - 100, - 20), - e.getMessage()); + Assert.assertTrue(e.getMessage().contains("100")); + Assert.assertTrue(e.getMessage().contains("5")); } finally { Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes()); buffer.close(); @@ -127,286 +100,6 @@ public void testMemoryAllocationMetrics() throws Exception { } } - @Test - public void testSmallFramesShrinkBufferWithoutGoingBelowInitialCapacity() throws Exception { - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - - try (AutoResizingBuffer buffer = new AutoResizingBuffer(100)) { - buffer.resizeIfNecessary(1600); - for (int size : new int[] {100, 4, 0}) { - int previousCapacity = buffer.array().length; - setLastShrinkTime(buffer, System.currentTimeMillis() - RpcUtils.MIN_SHRINK_INTERVAL - 1); - for (int i = 0; i <= RpcUtils.MAX_BUFFER_OVERSIZE_TIME; i++) { - buffer.resizeIfNecessary(size); - } - Assert.assertTrue(buffer.array().length < previousCapacity); - Assert.assertTrue(buffer.array().length >= 100); - Assert.assertEquals(buffer.array().length, memoryControl.getUsedMemoryInBytes()); - } - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testNewConnectionReclaimsIdleBuffersUnderMemoryPressure() throws Exception { - memoryControl.setLimit(80); - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - - byte[] payload = new byte[40]; - Arrays.fill(payload, (byte) 7); - try (TElasticFramedTransport transport = - new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { - transport.write(payload); - transport.flush(); - byte[] received = new byte[payload.length]; - transport.readAll(received, 0, received.length); - Assert.assertArrayEquals(payload, received); - Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); - - try (TElasticFramedTransport newConnection = - new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { - Assert.assertTrue(memoryControl.getUsedMemoryInBytes() <= 80); - Assert.assertTrue( - transport.readBuffer.getBuffer().length == 8 - || transport.writeBuffer.getBuffer().length == 8); - Assert.assertEquals(0, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); - transport.write(payload); - transport.flush(); - transport.readAll(received, 0, received.length); - Assert.assertArrayEquals(payload, received); - } - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test(timeout = 20000) - public void testThriftServerAcceptsNewConnectionAfterLargeFrame() throws Exception { - ServerSocket socket = new ServerSocket(); - socket.bind(new InetSocketAddress("127.0.0.1", 0)); - ExecutorService workers = Executors.newFixedThreadPool(2); - TServer server = - new TThreadPoolServer( - new TThreadPoolServer.Args(new TServerSocket(socket)) - .executorService(workers) - .transportFactory(new TElasticFramedTransport.Factory(8, 1024, true)) - .protocolFactory(new TBinaryProtocol.Factory()) - .processor( - (input, output) -> { - output.writeBinary(input.readBinary()); - output.getTransport().flush(); - })); - Thread serverThread = new Thread(server::serve, "rpc-buffer-memory-test-server"); - serverThread.setDaemon(true); - - // Client buffers are outside the server's memory control, as in a real deployment. - try (TElasticFramedTransport first = - new TElasticFramedTransport( - new TSocket("127.0.0.1", socket.getLocalPort(), 5000), 8, 1024, true); - TElasticFramedTransport second = - new TElasticFramedTransport( - new TSocket("127.0.0.1", socket.getLocalPort(), 5000), 8, 1024, true)) { - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - serverThread.start(); - first.open(); - TProtocol firstProtocol = new TBinaryProtocol(first); - byte[] payload = new byte[40]; - Arrays.fill(payload, (byte) 13); - firstProtocol.writeBinary(ByteBuffer.wrap(payload)); - first.flush(); - Assert.assertEquals(ByteBuffer.wrap(payload), firstProtocol.readBinary()); - - memoryControl.setLimit(memoryControl.getUsedMemoryInBytes()); - second.open(); - TProtocol secondProtocol = new TBinaryProtocol(second); - secondProtocol.writeBinary(ByteBuffer.wrap(new byte[] {1, 2, 3, 4})); - second.flush(); - Assert.assertEquals(ByteBuffer.wrap(new byte[] {1, 2, 3, 4}), secondProtocol.readBinary()); - Assert.assertEquals(0, AutoResizingBufferMemoryManager.getMemoryAllocationFailureCount()); - } finally { - server.stop(); - workers.shutdownNow(); - Assert.assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); - serverThread.join(5000); - Assert.assertFalse(serverThread.isAlive()); - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testMemoryPressureDoesNotReclaimUnflushedWriteBuffer() throws Exception { - memoryControl.setLimit(48); - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - - byte[] payload = new byte[40]; - Arrays.fill(payload, (byte) 9); - TMemoryBuffer wire = new TMemoryBuffer(0); - try (TElasticFramedTransport transport = new TElasticFramedTransport(wire, 8, 1024, true)) { - transport.write(payload); - Assert.assertThrows(IOException.class, () -> new AutoResizingBuffer(8)); - Assert.assertEquals(48, memoryControl.getUsedMemoryInBytes()); - - transport.flush(); - byte[] frame = new byte[4 + payload.length]; - wire.readAll(frame, 0, frame.length); - Assert.assertEquals(payload.length, TFramedTransport.decodeFrameSize(frame)); - Assert.assertArrayEquals(payload, Arrays.copyOfRange(frame, 4, frame.length)); - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testMemoryPressurePreservesUnreadFrame() throws Exception { - memoryControl.setLimit(80); - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - - byte[] payload = new byte[40]; - Arrays.fill(payload, (byte) 11); - try (TElasticFramedTransport transport = - new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, false)) { - transport.write(payload); - transport.flush(); - byte[] received = new byte[payload.length]; - transport.readAll(received, 0, 1); - Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); - - try (TElasticFramedTransport newConnection = - new TElasticFramedTransport(new TMemoryBuffer(0), 8, 1024, true)) { - Assert.assertEquals(40, transport.readBuffer.getBuffer().length); - Assert.assertEquals(8, transport.writeBuffer.getBuffer().length); - transport.readAll(received, 1, received.length - 1); - Assert.assertArrayEquals(payload, received); - } - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testSnappyConnectionReclaimsIdleCompressionBuffers() throws Exception { - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - - byte[] payload = new byte[4096]; - new Random(0).nextBytes(payload); - try (TSnappyElasticFramedTransport transport = - new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 32, 65536, true)) { - transport.write(payload); - transport.flush(); - byte[] received = new byte[payload.length]; - transport.readAll(received, 0, received.length); - Assert.assertArrayEquals(payload, received); - long retainedMemory = memoryControl.getUsedMemoryInBytes(); - memoryControl.setLimit(retainedMemory); - - try (TSnappyElasticFramedTransport newConnection = - new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 32, 65536, true)) { - Assert.assertTrue(memoryControl.getUsedMemoryInBytes() < retainedMemory); - memoryControl.setLimit(Long.MAX_VALUE); - transport.write(payload); - transport.flush(); - transport.readAll(received, 0, received.length); - Assert.assertArrayEquals(payload, received); - } - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testConcurrentAllocationsDoNotWaitForEachOthersBufferLocks() throws Exception { - memoryControl.setLimit(80); - CountDownLatch rejectedAllocations = new CountDownLatch(2); - AutoResizingBufferMemoryManager.setMemoryControl( - new AutoResizingBufferMemoryControl() { - @Override - public boolean allocate(long sizeInBytes) { - if (memoryControl.allocate(sizeInBytes)) { - return true; - } - // Both threads now hold their buffer locks while attempting to reclaim memory. - rejectedAllocations.countDown(); - try { - Assert.assertTrue(rejectedAllocations.await(5, TimeUnit.SECONDS)); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new AssertionError(e); - } - return false; - } - - @Override - public void release(long sizeInBytes) { - memoryControl.release(sizeInBytes); - } - }); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - ExecutorService executor = Executors.newFixedThreadPool(2); - try (AutoResizingBuffer first = new AutoResizingBuffer(8); - AutoResizingBuffer second = new AutoResizingBuffer(8)) { - first.resizeIfNecessary(40); - second.resizeIfNecessary(40); - Future firstFailure = - executor.submit( - () -> Assert.assertThrows(IOException.class, () -> first.resizeIfNecessary(41))); - Future secondFailure = - executor.submit( - () -> Assert.assertThrows(IOException.class, () -> second.resizeIfNecessary(41))); - Assert.assertNotNull(firstFailure.get(10, TimeUnit.SECONDS)); - Assert.assertNotNull(secondFailure.get(10, TimeUnit.SECONDS)); - Assert.assertEquals(80, memoryControl.getUsedMemoryInBytes()); - } finally { - executor.shutdownNow(); - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testMemoryPressureDoesNotReclaimBuffersFromAnotherControl() throws Exception { - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - try (AutoScalingBufferWriteTransport original = new AutoScalingBufferWriteTransport(8)) { - original.write(new byte[40]); - original.reset(); - - TestMemoryControl replacement = new TestMemoryControl(); - replacement.setLimit(0); - AutoResizingBufferMemoryManager.setMemoryControl(replacement); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - Assert.assertThrows(IOException.class, () -> new AutoResizingBuffer(8)); - Assert.assertEquals(40, original.getBuffer().length); - Assert.assertEquals(40, memoryControl.getUsedMemoryInBytes()); - Assert.assertEquals(0, replacement.getUsedMemoryInBytes()); - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - - @Test - public void testFailedOpenDoesNotLeakBufferMemoryWhenReconnecting() throws Exception { - memoryControl.setLimit(200); - AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); - AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); - - for (int i = 0; i < 10; i++) { - TSocket socket = - new TSocket("unavailable", 1) { - @Override - public void open() throws TTransportException { - throw new TTransportException("Simulated connection failure"); - } - }; - TElasticFramedTransport transport = new TElasticFramedTransport(socket, 50, 1024, true); - Assert.assertThrows(TTransportException.class, transport::open); - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - try (TElasticFramedTransport recovered = - new TElasticFramedTransport(new TMemoryBuffer(0), 50, 1024, true)) { - recovered.open(); - Assert.assertEquals(100, memoryControl.getUsedMemoryInBytes()); - } - Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); - } - @Test public void testElasticFramedTransportReleasesMemoryWhenClosed() throws Exception { AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); @@ -526,16 +219,10 @@ public void release(long sizeInBytes) { usedMemoryInBytes.addAndGet(-sizeInBytes); } - @Override - public long getUsedMemoryInBytes() { + private long getUsedMemoryInBytes() { return usedMemoryInBytes.get(); } - @Override - public long getTotalMemorySizeInBytes() { - return limit; - } - private void setLimit(long limit) { this.limit = limit; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/property/ClientPoolProperty.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/property/ClientPoolProperty.java index 1f818afe543da..a299cc056a361 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/property/ClientPoolProperty.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/property/ClientPoolProperty.java @@ -101,6 +101,9 @@ public ClientPoolProperty build() { poolConfig.setMaxIdlePerKey(maxIdleClientNumForEachNode); poolConfig.setTimeBetweenEvictionRuns(Duration.ofMillis(timeBetweenEvictionRuns)); poolConfig.setMinEvictableIdleTime(Duration.ofMillis(minIdleTimeForClient)); + // -1 checks all idle clients across all nodes in each run. The default of 3 can leave + // expired connections (and their RPC buffers) in a large pool for hours. + poolConfig.setNumTestsPerEvictionRun(-1); poolConfig.setMaxWait(Duration.ofMillis(waitClientTimeoutMs)); poolConfig.setTestOnReturn(true); poolConfig.setTestOnBorrow(true); @@ -114,7 +117,7 @@ private DefaultProperty() {} public static final long WAIT_CLIENT_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(30); public static final long MIN_IDLE_TIME_FOR_CLIENT_MS = TimeUnit.MINUTES.toMillis(1); - public static final long TIME_BETWEEN_EVICTION_RUNS_MS = TimeUnit.MINUTES.toMillis(1); + public static final long TIME_BETWEEN_EVICTION_RUNS_MS = TimeUnit.SECONDS.toMillis(30); public static final int MAX_CLIENT_NUM_FOR_EACH_NODE = 1000; public static final int MAX_IDLE_CLIENT_NUM_FOR_EACH_NODE = 1000; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java index 76fc2ba65b5ff..52b08556e3f02 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/memory/MemoryConfig.java @@ -64,16 +64,6 @@ public boolean allocate(long sizeInBytes) { public void release(long sizeInBytes) { memoryBlock.release(sizeInBytes); } - - @Override - public long getTotalMemorySizeInBytes() { - return memoryBlock.getTotalMemorySizeInBytes(); - } - - @Override - public long getUsedMemoryInBytes() { - return memoryBlock.getUsedMemoryInBytes(); - } }); } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/rpc/ClientPoolIdleEvictionTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/rpc/ClientPoolIdleEvictionTest.java new file mode 100644 index 0000000000000..dca5c7eb441ba --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/rpc/ClientPoolIdleEvictionTest.java @@ -0,0 +1,326 @@ +/* + * 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.iotdb.commons.client.property.ClientPoolProperty; + +import org.apache.commons.pool2.BaseKeyedPooledObjectFactory; +import org.apache.commons.pool2.PooledObject; +import org.apache.commons.pool2.impl.DefaultPooledObject; +import org.apache.commons.pool2.impl.GenericKeyedObjectPool; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.server.TServer; +import org.apache.thrift.server.TThreadPoolServer; +import org.apache.thrift.transport.TMemoryBuffer; +import org.apache.thrift.transport.TServerSocket; +import org.apache.thrift.transport.TSocket; +import org.apache.thrift.transport.TTransport; +import org.apache.thrift.transport.TTransportException; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.net.InetSocketAddress; +import java.net.ServerSocket; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +import static org.awaitility.Awaitility.await; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +@RunWith(Parameterized.class) +public class ClientPoolIdleEvictionTest { + + private static final int INITIAL_BUFFER_SIZE = 32; + private static final byte[] PAYLOAD = new byte[4096]; + + private final boolean compression; + private final TestMemoryControl memoryControl = new TestMemoryControl(); + + public ClientPoolIdleEvictionTest(boolean compression) { + this.compression = compression; + } + + @Parameterized.Parameters(name = "compression={0}") + public static Iterable parameters() { + return Arrays.asList(new Object[][] {{false}, {true}}); + } + + @Before + public void setUp() { + AutoResizingBufferMemoryManager.resetMemoryControl(); + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + AutoResizingBufferMemoryManager.setMemoryAllocateRetryIntervalInMs(1); + Arrays.fill(PAYLOAD, (byte) 13); + } + + @After + public void tearDown() { + AutoResizingBufferMemoryManager.resetMemoryControl(); + } + + @Test + public void testOneRunReclaimsAllExpiredClientsAcrossNodes() throws Exception { + try (GenericKeyedObjectPool pool = + createPool(100, -1, null)) { + TElasticFramedTransport active = pool.borrowObject(0); + long activeMemory = memoryControl.used.get(); + try { + List idleClients = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + idleClients.add(pool.borrowObject(i % 2)); + } + for (int i = 0; i < idleClients.size(); i++) { + pool.returnObject(i % 2, idleClients.get(i)); + } + + long retainedMemory = memoryControl.used.get(); + memoryControl.limit = retainedMemory; + assertThrows(TTransportException.class, () -> pool.borrowObject(2)); + assertEquals(retainedMemory, memoryControl.used.get()); + + // Wait for the idle timeout with the background evictor disabled, then run one sweep. + waitForIdleTimeout(); + pool.evict(); + + assertEquals(0, pool.getNumIdle()); + assertEquals(10, pool.getDestroyedCount()); + assertEquals(1, pool.getNumActive()); + assertEquals(activeMemory, memoryControl.used.get()); + + // Eviction must preserve a borrowed client's frame, and free quota for a new client. + byte[] actual = new byte[PAYLOAD.length]; + active.readAll(actual, 0, actual.length); + assertArrayEquals(PAYLOAD, actual); + TElasticFramedTransport replacement = pool.borrowObject(2); + pool.returnObject(2, replacement); + } finally { + pool.returnObject(0, active); + } + } + assertEquals(0, memoryControl.used.get()); + } + + @Test + public void testBackgroundEvictionReleasesIdleBufferMemory() throws Exception { + try (GenericKeyedObjectPool pool = + createPool(100, 20, null)) { + TElasticFramedTransport client = pool.borrowObject(0); + pool.returnObject(0, client); + + await() + .atMost(5, TimeUnit.SECONDS) + .pollInterval(Duration.ofMillis(10)) + .until(() -> pool.getNumIdle() == 0 && memoryControl.used.get() == 0); + assertEquals(1, pool.getDestroyedCount()); + } + assertEquals(0, memoryControl.used.get()); + } + + @Test + public void testUnexpiredClientRemainsReusable() throws Exception { + try (GenericKeyedObjectPool pool = + createPool(TimeUnit.MINUTES.toMillis(1), -1, null)) { + TElasticFramedTransport client = pool.borrowObject(0); + pool.returnObject(0, client); + long retainedMemory = memoryControl.used.get(); + + pool.evict(); + assertEquals(1, pool.getNumIdle()); + assertEquals(0, pool.getDestroyedCount()); + assertEquals(retainedMemory, memoryControl.used.get()); + + TElasticFramedTransport reused = pool.borrowObject(0); + try { + assertSame(client, reused); + byte[] actual = new byte[PAYLOAD.length]; + reused.readAll(actual, 0, actual.length); + assertArrayEquals(PAYLOAD, actual); + } finally { + pool.returnObject(0, reused); + } + } + assertEquals(0, memoryControl.used.get()); + } + + @Test(timeout = 20000) + public void testEvictionReleasesServerBuffersForNewConnections() throws Exception { + try (ServerSocket socket = new ServerSocket()) { + socket.bind(new InetSocketAddress("127.0.0.1", 0)); + ExecutorService workers = Executors.newFixedThreadPool(2); + TServer server = + new TThreadPoolServer( + new TThreadPoolServer.Args(new TServerSocket(socket)) + .executorService(workers) + .transportFactory( + compression + ? new TSnappyElasticFramedTransport.Factory( + INITIAL_BUFFER_SIZE, 65536, true) + : new TElasticFramedTransport.Factory(INITIAL_BUFFER_SIZE, 65536, true)) + .protocolFactory(new TBinaryProtocol.Factory()) + .processor( + (input, output) -> { + output.writeBinary(input.readBinary()); + output.getTransport().flush(); + })); + Thread serverThread = new Thread(server::serve, "client-pool-idle-eviction-test-server"); + serverThread.setDaemon(true); + serverThread.start(); + + try (GenericKeyedObjectPool pool = + createPool(100, -1, socket.getLocalPort())) { + TElasticFramedTransport first = pool.borrowObject(0); + try { + echo(first, PAYLOAD); + } finally { + pool.returnObject(0, first); + } + memoryControl.limit = memoryControl.used.get(); + assertThrows(TTransportException.class, () -> pool.borrowObject(1)); + + waitForIdleTimeout(); + pool.evict(); + assertFalse(first.isOpen()); + // The server must observe EOF and close its own read/write and compression buffers. + await() + .atMost(5, TimeUnit.SECONDS) + .pollInterval(Duration.ofMillis(10)) + .until(() -> memoryControl.used.get() == 0); + + TElasticFramedTransport replacement = pool.borrowObject(1); + try { + echo(replacement, new byte[] {1, 2, 3, 4}); + } finally { + pool.returnObject(1, replacement); + } + } finally { + server.stop(); + workers.shutdownNow(); + assertTrue(workers.awaitTermination(5, TimeUnit.SECONDS)); + serverThread.join(5000); + assertFalse(serverThread.isAlive()); + } + } + assertEquals(0, memoryControl.used.get()); + } + + private static void echo(TElasticFramedTransport transport, byte[] payload) throws Exception { + TProtocol protocol = new TBinaryProtocol(transport); + protocol.writeBinary(ByteBuffer.wrap(payload)); + transport.flush(); + assertEquals(ByteBuffer.wrap(payload), protocol.readBinary()); + } + + private static void waitForIdleTimeout() { + long returnedAt = System.nanoTime(); + await() + .atMost(5, TimeUnit.SECONDS) + .pollInterval(Duration.ofMillis(10)) + .until(() -> System.nanoTime() - returnedAt > TimeUnit.MILLISECONDS.toNanos(200)); + } + + private GenericKeyedObjectPool createPool( + long idleTimeoutMs, long evictionIntervalMs, Integer serverPort) { + return new GenericKeyedObjectPool<>( + new BaseKeyedPooledObjectFactory() { + @Override + public TElasticFramedTransport create(Integer node) throws Exception { + TElasticFramedTransport transport = + createTransport( + serverPort == null + ? new TMemoryBuffer(0) + : new TSocket("127.0.0.1", serverPort, 5000)); + try { + if (serverPort == null) { + transport.write(PAYLOAD, 0, PAYLOAD.length); + transport.flush(); + } else { + transport.open(); + } + return transport; + } catch (Exception e) { + transport.close(); + throw e; + } + } + + @Override + public PooledObject wrap(TElasticFramedTransport transport) { + return new DefaultPooledObject<>(transport); + } + + @Override + public void destroyObject(Integer node, PooledObject transport) { + transport.getObject().close(); + } + }, + new ClientPoolProperty.Builder() + .setMaxClientNumForEachNode(16) + .setMaxIdleClientNumForEachNode(16) + .setMinIdleTimeForClient(idleTimeoutMs) + .setTimeBetweenEvictionRuns(evictionIntervalMs) + .build() + .getConfig()); + } + + private TElasticFramedTransport createTransport(TTransport underlying) throws Exception { + return compression + ? new TSnappyElasticFramedTransport(underlying, INITIAL_BUFFER_SIZE, 65536, true) + : new TElasticFramedTransport(underlying, INITIAL_BUFFER_SIZE, 65536, true); + } + + private static class TestMemoryControl implements AutoResizingBufferMemoryControl { + private final AtomicLong used = new AtomicLong(); + private volatile long limit = Long.MAX_VALUE; + + @Override + public boolean allocate(long bytes) { + while (true) { + long current = used.get(); + if (current + bytes > limit) { + return false; + } + if (used.compareAndSet(current, current + bytes)) { + return true; + } + } + } + + @Override + public void release(long bytes) { + used.addAndGet(-bytes); + } + } +}