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 cd925b4863af1..fe70f9cbe3cf4 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 @@ -127,4 +127,25 @@ private RpcMessages() {} public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 = "Data type %s is not supported."; + public static final String EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA = + "Invalid buffer size: %d"; + + public static final String EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D = + "Write buffer size (%d) larger than maximum buffer size (%d)!"; + + public static final String EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D = + "Invalid compression buffer size: %d"; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA = + "Serialized frame size is at least %d bytes, exceeding the maximum Java array size of %d bytes."; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9 = + "Serialized frame size is at least %d bytes, exceeding the configured Thrift max frame size of %d bytes."; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C = + "Serialized frame size is at least %d bytes, exceeding the maximum Java array size of %d bytes to %s."; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898 = + "Serialized frame size is at least %d bytes, exceeding the configured Thrift max frame size of %d bytes to %s."; + } 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 fc059e26f5f32..d4b26bba60b2a 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 @@ -108,4 +108,25 @@ private RpcMessages() {} "通过列名 %s 获取的值为 NULL。"; public static final String EXCEPTION_DATA_TYPE_ARG_IS_NOT_SUPPORTED_31213160 = "不支持数据类型 %s。"; + public static final String EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA = + "无效的缓冲区大小：%d"; + + public static final String EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D = + "写缓冲区大小（%d）超过最大缓冲区大小（%d）！"; + + public static final String EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D = + "无效的压缩缓冲区大小：%d"; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA = + "序列化帧大小至少为 %d 字节，超过 Java 数组最大大小 %d 字节。"; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9 = + "序列化帧大小至少为 %d 字节，超过配置的 Thrift 最大帧大小 %d 字节。"; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C = + "序列化帧大小至少为 %d 字节，超过 Java 数组最大大小 %d 字节，目标地址为 %s。"; + + public static final String EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898 = + "序列化帧大小至少为 %d 字节，超过配置的 Thrift 最大帧大小 %d 字节，目标地址为 %s。"; + } 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 f2a57cb5ca862..26c5eb97779d9 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 @@ -19,6 +19,8 @@ package org.apache.iotdb.rpc; +import org.apache.iotdb.rpc.i18n.RpcMessages; + import java.io.IOException; import java.util.Arrays; @@ -35,6 +37,8 @@ */ class AutoResizingBuffer implements AutoCloseable { + static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8; + private byte[] array; private int bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; private final int initialCapacity; @@ -71,21 +75,38 @@ public void resizeIfNecessary(int size) throws IOException { } void ensureCapacity(int size) throws IOException { + ensureCapacity(size, MAX_ARRAY_SIZE); + } + + void ensureCapacity(int size, int maxCapacity) throws IOException { + int limit = Math.min(maxCapacity, MAX_ARRAY_SIZE); + if (size < 0 || size > limit) { + throw new IllegalArgumentException( + String.format(RpcMessages.EXCEPTION_INVALID_BUFFER_SIZE_ARG_5131C6CA, size)); + } reserveCurrentCapacityIfReleased(); final int currentCapacity = this.array.length; if (currentCapacity < size) { - // Increase by a factor of 1.5x - int growCapacity = currentCapacity + (currentCapacity >> 1); - int newCapacity = Math.max(growCapacity, size); - resize(newCapacity, size); + resize(Math.min(calculateNewCapacity(currentCapacity, size), limit), size); bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; } } + static int calculateNewCapacity(int currentCapacity, int requiredCapacity) { + // Keep geometric growth even above 1.33 GiB. An int overflow here would otherwise copy the + // entire buffer for every small write during Thrift serialization. + long growCapacity = currentCapacity + (long) (currentCapacity >> 1); + return (int) Math.min(MAX_ARRAY_SIZE, Math.max(growCapacity, requiredCapacity)); + } + public byte[] array() { return this.array; } + boolean isClosed() { + return closed; + } + void resetToInitialCapacity() throws IOException { reserveCurrentCapacityIfReleased(); resize(initialCapacity, initialCapacity); 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 9a71a1fa8865d..4450d8b6b407f 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 @@ -19,6 +19,8 @@ package org.apache.iotdb.rpc; +import org.apache.iotdb.rpc.i18n.RpcMessages; + import org.apache.thrift.TConfiguration; import org.apache.thrift.transport.TTransportException; @@ -55,10 +57,25 @@ public int read(byte[] buf, int off, int len) { @Override public void write(byte[] toWrite, int off, int len) throws TTransportException { + write(toWrite, off, len, AutoResizingBuffer.MAX_ARRAY_SIZE); + } + + void write(byte[] toWrite, int off, int len, int maxCapacity) throws TTransportException { + int limit = Math.min(maxCapacity, AutoResizingBuffer.MAX_ARRAY_SIZE); + long requiredCapacity = (long) pos + len; + if (requiredCapacity > limit) { + throw new TTransportException( + TTransportException.CORRUPTED_DATA, + String.format( + RpcMessages + .EXCEPTION_WRITE_BUFFER_SIZE_ARG_LARGER_THAN_MAXIMUM_BUFFER_SIZE_ARG_C0A47F3D, + requiredCapacity, + limit)); + } + // A small prefix during serialization does not indicate a small frame. Only the complete + // frame size, supplied after flushing, may trigger shrinking. try { - // A small prefix during serialization does not indicate a small frame. Only the complete - // frame size, supplied after flushing, may trigger shrinking. - ensureCapacity(pos + len); + ensureCapacity((int) requiredCapacity, limit); } catch (IOException e) { throw new TTransportException(e); } @@ -66,6 +83,22 @@ public void write(byte[] toWrite, int off, int len) throws TTransportException { pos += len; } + // Framed callers keep capacity within their limit. A fitting write needs no further limit + // checks; compare remaining space rather than adding pos and len, which could overflow. + boolean tryWriteWithinCapacity(byte[] toWrite, int off, int len) { + byte[] array = getBuffer(); + // A closed reusable buffer must reserve its memory again. Crossing the reuse threshold + // must also go through ensureCapacity so the large frame uses temporary storage. + if (buf.isClosed() + || len > array.length - pos + || (temporaryBuffer == null && len > maxReusableBufferSize - pos)) { + return false; + } + System.arraycopy(toWrite, off, array, pos, len); + pos += len; + return true; + } + public int getPos() { return pos; } @@ -101,6 +134,10 @@ public void resizeIfNecessary(int size) throws IOException { } private void ensureCapacity(int size) throws IOException { + ensureCapacity(size, AutoResizingBuffer.MAX_ARRAY_SIZE); + } + + private void ensureCapacity(int size, int maxCapacity) throws IOException { if (size > maxReusableBufferSize) { if (temporaryBuffer == null || temporaryBuffer.length < size) { // Preserve an unfinished frame when incremental serialization crosses the threshold. @@ -108,13 +145,15 @@ private void ensureCapacity(int size) throws IOException { int capacity = temporaryBuffer == null ? size - : Math.max(size, temporaryBuffer.length + (temporaryBuffer.length >> 1)); + : Math.min( + AutoResizingBuffer.calculateNewCapacity(temporaryBuffer.length, size), + maxCapacity); byte[] newBuffer = Arrays.copyOf(getBuffer(), capacity); buf.resetToInitialCapacity(); temporaryBuffer = newBuffer; } } else { - buf.ensureCapacity(size); + buf.ensureCapacity(size, maxCapacity); } } 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 ce4440a9e1042..aa03783f81e9d 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 @@ -19,6 +19,8 @@ package org.apache.iotdb.rpc; +import org.apache.iotdb.rpc.i18n.RpcMessages; + import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.apache.thrift.transport.layered.TFramedTransport; @@ -75,6 +77,13 @@ protected void closeAllocatedBuffers() { } } + @Override + protected int getMaxWriteBufferSize() { + // Preserve the wire-size limit: a compressible payload may exceed thriftMaxFrameSize before + // compression. Incremental writes still must fit in a Java array. + return AutoResizingBuffer.MAX_ARRAY_SIZE; + } + @Override protected void readFrame() throws TTransportException { // Discard the previous frame's budget before reading the next frame header and payload. @@ -88,6 +97,7 @@ protected void readFrame() throws TTransportException { RpcStat.readCompressedBytes.addAndGet(size); try { int uncompressedLength = uncompressedLength(readBuffer.getBuffer(), 0, size); + checkBufferSize(uncompressedLength); RpcStat.readBytes.addAndGet(uncompressedLength); readCompressBuffer.resizeIfNecessary(uncompressedLength); uncompress(readBuffer.getBuffer(), 0, size, readCompressBuffer.getBuffer(), 0); @@ -105,6 +115,7 @@ public void flush() throws TTransportException { RpcStat.writeBytes.addAndGet(length); try { int maxCompressedLength = maxCompressedLength(length); + checkBufferSize(maxCompressedLength); writeCompressBuffer.resizeIfNecessary(maxCompressedLength); int compressedLength = compress(writeBuffer.getBuffer(), 0, length, writeCompressBuffer.getBuffer(), 0); @@ -128,6 +139,17 @@ public void flush() throws TTransportException { underlying.flush(); } + private void checkBufferSize(int size) throws TTransportException { + if (size < 0 || size > AutoResizingBuffer.MAX_ARRAY_SIZE) { + // Compression bounds can overflow int even when the input fits in a Java array. Reject + // before resizing and invalidate the transport, just as for an oversized wire frame. + close(); + throw new TTransportException( + TTransportException.CORRUPTED_DATA, + String.format(RpcMessages.EXCEPTION_INVALID_COMPRESSION_BUFFER_SIZE_ARG_D677543D, size)); + } + } + protected abstract int uncompressedLength(byte[] but, int off, int len) throws IOException; protected abstract int maxCompressedLength(int len); 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 8428ef3c94a4e..7079b7987ea47 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 @@ -43,9 +43,8 @@ public class TElasticFramedTransport extends TTransport { public static class Factory extends TTransportFactory { /** - * It is used to prevent the size of the parsing package from being too large and allocating the - * buffer will cause oom. Therefore, the maximum length of the requested memory is limited when - * reading. + * Maximum size of a frame on the wire. Uncompressed writes are checked before buffering; + * compressed writes are checked after compression. */ protected final int thriftMaxFrameSize; @@ -106,7 +105,11 @@ protected TElasticFramedTransport( readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize, maxReusableBufferSize); writeBuffer = - new AutoScalingBufferWriteTransport(thriftDefaultBufferSize, maxReusableBufferSize); + new AutoScalingBufferWriteTransport( + Math.min( + thriftDefaultBufferSize, + Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE)), + maxReusableBufferSize); } catch (IOException e) { closeAllocatedBuffers(); throw new TTransportException(e); @@ -383,7 +386,50 @@ public void resetMessageSizeAndConsumedBytes(long newSize) throws TTransportExce @Override public void write(byte[] buf, int off, int len) throws TTransportException { - writeBuffer.write(buf, off, len); + if (writeBuffer.tryWriteWithinCapacity(buf, off, len)) { + return; + } + long size = (long) writeBuffer.getPos() + len; + int maxSize = getMaxWriteBufferSize(); + if (size > maxSize) { + // Thrift serializes the whole request before flush. Reject it before growing the buffer so + // an oversized request cannot keep allocating and retain the caller's query memory. + // TSocket.close() clears its socket, so retain the destination before invalidating it. + SocketAddress remoteAddress = null; + if (underlying instanceof TSocket socket && socket.getSocket() != null) { + remoteAddress = socket.getSocket().getRemoteSocketAddress(); + } + close(); + String message; + if (remoteAddress == null) { + message = + String.format( + maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE + ? RpcMessages + .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_E21E05CA + : RpcMessages + .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_1CDD4BB9, + size, + maxSize); + } else { + message = + String.format( + maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE + ? RpcMessages + .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_MAXIMUM_JAVA_ARRAY_SIZE_OF_ARG_BYTES_TO_ARG_6F166F0C + : RpcMessages + .EXCEPTION_SERIALIZED_FRAME_SIZE_IS_AT_LEAST_ARG_BYTES_EXCEEDING_THE_CONFIGURED_THRIFT_MAX_FRAME_SIZE_OF_ARG_BYTES_TO_ARG_7EF64898, + size, + maxSize, + remoteAddress); + } + throw new TTransportException(TTransportException.CORRUPTED_DATA, message); + } + writeBuffer.write(buf, off, len, maxSize); + } + + protected int getMaxWriteBufferSize() { + return Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE); } public TTransport getSocket() { 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 18aed44eec3cf..3ffffda811bb1 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 @@ -40,6 +40,137 @@ public void tearDown() { AutoResizingBufferMemoryManager.resetMemoryControl(); } + @Test + public void testGrowthNearIntegerLimitKeepsHeadroom() { + Assert.assertEquals( + AutoResizingBuffer.MAX_ARRAY_SIZE, + AutoResizingBuffer.calculateNewCapacity(1_500_000_000, 1_500_000_008)); + Assert.assertEquals( + AutoResizingBuffer.MAX_ARRAY_SIZE, + AutoResizingBuffer.calculateNewCapacity( + AutoResizingBuffer.MAX_ARRAY_SIZE - 8, AutoResizingBuffer.MAX_ARRAY_SIZE)); + } + + @Test + public void testInvalidBufferSizeDoesNotResize() throws Exception { + AutoResizingBuffer buffer = new AutoResizingBuffer(8); + for (int size : new int[] {-1, Integer.MAX_VALUE}) { + try { + buffer.ensureCapacity(size); + Assert.fail("The requested capacity must fit in a Java array"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(8, buffer.array().length); + } + } + } + + @Test + public void testWritePositionOverflowIsRejectedBeforeCopying() throws Exception { + AutoScalingBufferWriteTransport transport = new AutoScalingBufferWriteTransport(8); + Field position = AutoScalingBufferWriteTransport.class.getDeclaredField("pos"); + position.setAccessible(true); + position.setInt(transport, Integer.MAX_VALUE - 4); + byte[] buffer = transport.getBuffer(); + try { + transport.write(new byte[8]); + Assert.fail("The write position must not overflow"); + } catch (TTransportException e) { + Assert.assertEquals(TTransportException.CORRUPTED_DATA, e.getType()); + Assert.assertSame(buffer, transport.getBuffer()); + Assert.assertEquals(Integer.MAX_VALUE - 4, transport.getPos()); + } + } + + @Test + public void testGrowthPreservesDataAndHeadroom() throws Exception { + AutoResizingBuffer buffer = new AutoResizingBuffer(100); + buffer.array()[99] = 42; + buffer.resizeIfNecessary(110); + Assert.assertEquals(150, buffer.array().length); + Assert.assertEquals(42, buffer.array()[99]); + + buffer.resizeIfNecessary(1000); + Assert.assertEquals(1000, buffer.array().length); + Assert.assertEquals(42, buffer.array()[99]); + } + + @Test + public void testFramedFastPathReservesMemoryAfterReopening() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 32, 128, false)) { + Assert.assertEquals(64, memoryControl.getUsedMemoryInBytes()); + transport.write(new byte[8]); + transport.close(); + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + + transport.open(); + transport.write(new byte[8]); + Assert.assertEquals(32, memoryControl.getUsedMemoryInBytes()); + Assert.assertEquals(8, transport.writeBuffer.getPos()); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testCrossingReuseThresholdWithinCapacityUsesTemporaryBuffer() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 100, false, 32)) { + byte[] payload = new byte[36]; + new Random(0).nextBytes(payload); + for (int offset = 0; offset < 32; offset += 8) { + transport.write(payload, offset, 8); + } + byte[] reusableBuffer = transport.writeBuffer.getBuffer(); + Assert.assertEquals(36, reusableBuffer.length); + Assert.assertEquals(44, memoryControl.getUsedMemoryInBytes()); + + // This write crosses the reuse threshold while still fitting in the reusable array. + transport.write(payload, 32, 4); + Assert.assertNotSame(reusableBuffer, transport.writeBuffer.getBuffer()); + Assert.assertArrayEquals(payload, transport.writeBuffer.getBuffer()); + Assert.assertEquals(16, memoryControl.getUsedMemoryInBytes()); + transport.flush(); + Assert.assertEquals(8, transport.writeBuffer.getBuffer().length); + Assert.assertEquals(16, memoryControl.getUsedMemoryInBytes()); + + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + + @Test + public void testTemporaryWriteBufferGrowthRespectsFrameLimit() throws Exception { + AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 100, false, 32)) { + byte[] payload = new byte[100]; + new Random(0).nextBytes(payload); + for (int frame = 0; frame < 2; frame++) { + for (int offset = 0; offset < payload.length; offset += 8) { + transport.write(payload, offset, Math.min(8, payload.length - offset)); + Assert.assertTrue(transport.writeBuffer.getBuffer().length <= payload.length); + } + Assert.assertEquals(100, transport.writeBuffer.getBuffer().length); + transport.flush(); + Assert.assertEquals(8, transport.writeBuffer.getBuffer().length); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + Assert.assertArrayEquals(payload, actual); + } + transport.write(payload); + TTransportException failure = + Assert.assertThrows(TTransportException.class, () -> transport.write(new byte[1])); + Assert.assertEquals(TTransportException.CORRUPTED_DATA, failure.getType()); + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + Assert.assertEquals(8, transport.writeBuffer.getBuffer().length); + } + Assert.assertEquals(0, memoryControl.getUsedMemoryInBytes()); + } + @Test public void testAllocateAndReleaseMemoryWhenResizing() throws Exception { AutoResizingBufferMemoryManager.setMemoryControl(memoryControl); diff --git a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java index 0cbab4409e585..4b7fefa25da46 100644 --- a/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java +++ b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java @@ -22,19 +22,271 @@ import org.apache.thrift.TConfiguration; import org.apache.thrift.transport.TByteBuffer; import org.apache.thrift.transport.TMemoryBuffer; +import org.apache.thrift.transport.TSocket; import org.apache.thrift.transport.TTransportException; import org.junit.Test; +import java.lang.reflect.Field; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.Arrays; +import java.util.Random; 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; import static org.junit.Assert.fail; public class TElasticFramedTransportTest { + @Test + public void testEarlyFrameLimitErrorKeepsRemoteEndpoint() throws Exception { + assertEarlyWriteRejectionKeepsRemoteEndpoint(false); + } + + @Test + public void testEarlyArrayLimitErrorKeepsRemoteEndpoint() throws Exception { + assertEarlyWriteRejectionKeepsRemoteEndpoint(true); + } + + private void assertEarlyWriteRejectionKeepsRemoteEndpoint(boolean arrayLimit) throws Exception { + try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress()); + Socket client = new Socket(server.getInetAddress(), server.getLocalPort()); + Socket accepted = server.accept(); + TSocket underlying = new TSocket(client); + TElasticFramedTransport transport = + arrayLimit + ? new TSnappyElasticFramedTransport(underlying, 8, 64, false) + : new TElasticFramedTransport(underlying, 8, 64, false)) { + String remoteEndpoint = client.getRemoteSocketAddress().toString(); + if (arrayLimit) { + Field position = AutoScalingBufferWriteTransport.class.getDeclaredField("pos"); + position.setAccessible(true); + position.setInt(transport.writeBuffer, AutoResizingBuffer.MAX_ARRAY_SIZE - 4); + } else { + transport.write(new byte[64]); + } + + TTransportException failure = + assertThrows(TTransportException.class, () -> transport.write(new byte[16])); + assertEquals(TTransportException.CORRUPTED_DATA, failure.getType()); + assertTrue(failure.getMessage(), failure.getMessage().contains(" to " + remoteEndpoint)); + assertTrue(client.isClosed()); + accepted.setSoTimeout(1000); + assertEquals(-1, accepted.getInputStream().read()); + } + } + + @Test + public void testRejectOversizedFrameBeforeGrowingBuffer() throws Exception { + for (int prefixLength : new int[] {0, 64}) { + boolean[] closed = {false}; + TMemoryBuffer underlying = + new TMemoryBuffer(0) { + @Override + public void close() { + closed[0] = true; + super.close(); + } + }; + TElasticFramedTransport transport = new TElasticFramedTransport(underlying, 8, 64, false); + transport.write(new byte[prefixLength]); + byte[] buffer = transport.writeBuffer.getBuffer(); + try { + transport.write(new byte[65 - prefixLength]); + fail("An oversized frame must fail during serialization, before flush"); + } catch (TTransportException e) { + assertEquals(TTransportException.CORRUPTED_DATA, e.getType()); + assertTrue(e.getMessage(), e.getMessage().contains("at least 65 bytes")); + assertTrue( + e.getMessage(), + e.getMessage().contains("configured Thrift max frame size of 64 bytes")); + assertFalse(e.getMessage(), e.getMessage().contains("dn_thrift_max_frame_size")); + assertSame(buffer, transport.writeBuffer.getBuffer()); + assertEquals(0, transport.writeBuffer.getPos()); + assertEquals(0, underlying.length()); + assertTrue(closed[0]); + } + } + } + + @Test + public void testIncrementalGrowthDoesNotExceedFrameLimit() throws Exception { + int limit = 1024 * 1024; + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 1024, limit, false)) { + byte[] chunk = new byte[8]; + Arrays.fill(chunk, (byte) 42); + for (int written = 0; written < limit; written += chunk.length) { + transport.write(chunk); + assertTrue( + "Allocated " + transport.writeBuffer.getBuffer().length + " bytes for limit " + limit, + transport.writeBuffer.getBuffer().length <= limit); + } + assertEquals(limit, transport.writeBuffer.getBuffer().length); + assertEquals(42, transport.writeBuffer.getBuffer()[0]); + assertEquals(42, transport.writeBuffer.getBuffer()[limit - 1]); + TTransportException failure = + assertThrows(TTransportException.class, () -> transport.write(chunk)); + assertTrue(failure.getMessage(), failure.getMessage().contains("at least 1048584 bytes")); + assertEquals(limit, transport.writeBuffer.getBuffer().length); + } + } + + @Test + public void testInitialWriteBufferRespectsFrameLimit() throws Exception { + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 1024, 64, false)) { + assertEquals(64, transport.writeBuffer.getBuffer().length); + byte[] payload = new byte[64]; + payload[63] = 42; + for (int frame = 0; frame < 2; frame++) { + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + assertEquals(64, transport.writeBuffer.getBuffer().length); + } + } + } + + @Test + public void testFrameAtSizeLimit() throws Exception { + TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 8, 64, false); + byte[] payload = new byte[64]; + payload[63] = 42; + for (int i = 0; i < payload.length; i += 8) { + transport.write(payload, i, 8); + } + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + } + + @Test + public void testCompressedFrameLimitAppliesAfterCompression() throws Exception { + TElasticFramedTransport transport = + new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 128, 64, false); + byte[] payload = new byte[256]; + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + } + + @Test + public void testCompressedFrameLimitIsCheckedAfterBufferedWrites() throws Exception { + boolean[] closed = {false}; + TMemoryBuffer underlying = + new TMemoryBuffer(0) { + @Override + public void close() { + closed[0] = true; + super.close(); + } + }; + TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport(underlying, 8, 64, false); + byte[] payload = new byte[256]; + transport.write(payload); + transport.flush(); + byte[] actual = new byte[payload.length]; + transport.readAll(actual, 0, actual.length); + assertArrayEquals(payload, actual); + byte[] buffer = transport.writeBuffer.getBuffer(); + + new Random(0).nextBytes(payload); + transport.write(payload); + assertSame(buffer, transport.writeBuffer.getBuffer()); + TTransportException failure = assertThrows(TTransportException.class, transport::flush); + assertEquals(TTransportException.CORRUPTED_DATA, failure.getType()); + assertTrue(closed[0]); + } + + @Test + public void testBufferedWriteCannotBypassArrayLimitOnOverflow() throws Exception { + for (boolean compressed : new boolean[] {false, true}) { + boolean[] closed = {false}; + TMemoryBuffer underlying = + new TMemoryBuffer(0) { + @Override + public void close() { + closed[0] = true; + super.close(); + } + }; + TElasticFramedTransport transport = + compressed + ? new TSnappyElasticFramedTransport(underlying, 8, 64, false) + : new TElasticFramedTransport( + underlying, 8, AutoResizingBuffer.MAX_ARRAY_SIZE, false); + Field position = AutoScalingBufferWriteTransport.class.getDeclaredField("pos"); + position.setAccessible(true); + int oldPosition = AutoResizingBuffer.MAX_ARRAY_SIZE - 4; + position.setInt(transport.writeBuffer, oldPosition); + final byte[] buffer = transport.writeBuffer.getBuffer(); + TTransportException failure = + assertThrows(TTransportException.class, () -> transport.write(new byte[16])); + assertEquals(TTransportException.CORRUPTED_DATA, failure.getType()); + assertTrue(failure.getMessage(), failure.getMessage().contains("maximum Java array size")); + assertFalse(failure.getMessage(), failure.getMessage().contains("dn_thrift_max_frame_size")); + assertSame(buffer, transport.writeBuffer.getBuffer()); + assertEquals(0, transport.writeBuffer.getPos()); + assertTrue(closed[0]); + } + } + + @Test + public void testInvalidCompressionBufferSizesCloseTransport() throws Exception { + for (int size : new int[] {-1, Integer.MAX_VALUE}) { + for (boolean read : new boolean[] {false, true}) { + boolean[] closed = {false}; + TMemoryBuffer underlying = + new TMemoryBuffer(0) { + @Override + public void close() { + closed[0] = true; + super.close(); + } + }; + TSnappyElasticFramedTransport transport = + new TSnappyElasticFramedTransport(underlying, 8, 64, false) { + @Override + protected int maxCompressedLength(int length) { + // Snappy's bound overflows int for large, otherwise valid input lengths. + return size; + } + + @Override + protected int uncompressedLength(byte[] bytes, int offset, int length) { + return size; + } + }; + TTransportException failure; + if (read) { + underlying.write(new byte[] {0, 0, 0, 1, 0}); + failure = + assertThrows(TTransportException.class, () -> transport.read(new byte[1], 0, 1)); + } else { + transport.write(new byte[1]); + failure = assertThrows(TTransportException.class, transport::flush); + } + assertEquals(TTransportException.CORRUPTED_DATA, failure.getType()); + assertTrue(closed[0]); + } + } + } + @Test public void testReadFramesResetMessageSize() throws TTransportException { byte[] firstFrame = {1, 2}; diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 7b627c78386da..378a6381dcc99 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -4974,4 +4974,13 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ARG_ARG_MS_ARG_INVOCATIONS_ARG_APPLICATIONS_F056BD84 = "%s: %s ms, %s invocations, %s applications"; + public static final String EXCEPTION_NOT_ENOUGH_MEMORY_WHILE_ANALYZING_TIME_PARTITIONS_NARROW_THE_QUERY_TIME_RANGE_OR_INCREASE_QUERY_MEMORY_MEMORY_DETAILS_REQUESTED_THIS_TIME_ARG_CURRENT_FREE_MEMORY_ARG_ORIGINAL_ERROR_ARG_37676B2A = + "Not enough memory while analyzing time partitions. Narrow the query time range or increase query memory. Memory details: requested this time %s, current free memory %s. Original error: %s"; + + public static final String EXCEPTION_QUERY_ARG_CANNOT_RESERVE_ARG_BYTES_PENDING_FRONT_END_MEMORY_RESERVATION_EXCEEDS_LONG_MAX_VALUE_BYTES_D38CD6D8 = + "Query %s cannot reserve %s bytes: pending front-end memory reservation exceeds Long.MAX_VALUE bytes."; + + public static final String LOG_QUERYFAILED_QUERYID_ARG_SQL_ARG_249264E9 = + "[QueryFailed] queryId: {}, sql: {}"; + } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index e710c42235486..02cf8edc0a469 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -5675,4 +5675,13 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ARG_ARG_MS_ARG_INVOCATIONS_ARG_APPLICATIONS_F056BD84 = "%s：%s ms，调用 %s 次，应用 %s 次"; + public static final String EXCEPTION_NOT_ENOUGH_MEMORY_WHILE_ANALYZING_TIME_PARTITIONS_NARROW_THE_QUERY_TIME_RANGE_OR_INCREASE_QUERY_MEMORY_MEMORY_DETAILS_REQUESTED_THIS_TIME_ARG_CURRENT_FREE_MEMORY_ARG_ORIGINAL_ERROR_ARG_37676B2A = + "分析时间分区时内存不足。请缩小查询时间范围或增加查询内存。内存详情：本次请求 %s，当前空闲内存 %s。原始错误：%s"; + + public static final String EXCEPTION_QUERY_ARG_CANNOT_RESERVE_ARG_BYTES_PENDING_FRONT_END_MEMORY_RESERVATION_EXCEEDS_LONG_MAX_VALUE_BYTES_D38CD6D8 = + "查询 %s 无法预留 %s 字节：待预留的前端内存超过 Long.MAX_VALUE 字节。"; + + public static final String LOG_QUERYFAILED_QUERYID_ARG_SQL_ARG_249264E9 = + "[QueryFailed] queryId: {}, SQL：{}"; + } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java index ad9c932164cbe..e2b84a0960b61 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java @@ -700,6 +700,35 @@ public void reserveMemoryForFrontEndImmediately() { } } + public void reserveMemoryForTimePartitionSlots(final long bytes) { + try { + this.memoryReservationManager.reserveMemoryCumulatively(bytes); + } catch (MemoryNotEnoughException e) { + long requestedBytes = bytes; + long freeBytes; + if (e instanceof OperatorMemoryNotEnoughException memoryFailure) { + requestedBytes = memoryFailure.getRequestedBytes(); + freeBytes = memoryFailure.getFreeBytes(); + } else { + freeBytes = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + } + throw new MemoryNotEnoughException( + String.format( + Locale.ROOT, + DataNodeQueryMessages + .EXCEPTION_NOT_ENOUGH_MEMORY_WHILE_ANALYZING_TIME_PARTITIONS_NARROW_THE_QUERY_TIME_RANGE_OR_INCREASE_QUERY_MEMORY_MEMORY_DETAILS_REQUESTED_THIS_TIME_ARG_CURRENT_FREE_MEMORY_ARG_ORIGINAL_ERROR_ARG_37676B2A, + formatBytes(requestedBytes), + formatBytes(freeBytes), + e.getMessage())); + } + } + + public void finishResultSetColumnMemoryTracking() { + // Column expansion has finished. Later partition/planning reservations must not be diagnosed + // as excessive result columns just because this context still contains column statistics. + resultSetColumnMemoryTrackingEnabled = false; + } + public void releaseAllMemoryReservedForFrontEnd() { this.memoryReservationManager.releaseAllReservedMemory(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java index eda408bd7f59d..32ea59c26c20d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java @@ -320,6 +320,8 @@ private ExecutionResult execution( return execution(queryId, session, sql, userQuery, debug, false, iQueryExecutionFactory); } + // Record the SQL for Errors as well, preserving the original failure and finally cleanup. + @SuppressWarnings("java:S1181") private ExecutionResult execution( long queryId, SessionInfo session, @@ -369,6 +371,13 @@ private ExecutionResult execution( result.status.setNeedRetry(true); } return result; + } catch (RuntimeException | Error e) { + // Analysis can fail before registration. Only link the query ID to its SQL here; callers + // classify the failure and log its stack, avoiding duplicate stacks and WARNs for user + // errors. + LOGGER.info( + DataNodeQueryMessages.LOG_QUERYFAILED_QUERYID_ARG_SQL_ARG_249264E9, globalQueryId, sql); + throw e; } finally { if (queryContext != null) { queryContext.releaseAllMemoryReservedForFrontEnd(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java index 7aca24384e239..c40c3e57f985a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java @@ -49,6 +49,7 @@ import org.apache.iotdb.db.exception.ainode.GetModelInfoException; import org.apache.iotdb.db.exception.metadata.template.TemplateIncompatibleException; import org.apache.iotdb.db.exception.metadata.view.UnsupportedViewException; +import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.DeviceContext; @@ -2276,6 +2277,8 @@ private DataPartition fetchDataPartitionByDevices( @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning public static Pair, Pair> getTimePartitionSlotList( Filter timeFilter, MPPQueryContext context) { + checkTimePartitionQueryTimeout(context); + context.finishResultSetColumnMemoryTracking(); if (timeFilter == null) { // (-oo, +oo) return new Pair<>(Collections.emptyList(), new Pair<>(true, true)); @@ -2319,7 +2322,13 @@ public static Pair, Pair> getTimePart TimeRange currentTimeRange = timeRangeList.get(index); reserveMemoryForTimePartitionSlot( currentTimeRange.getMax(), currentTimeRange.getMin(), context); + int iterations = 0; while (index < size) { + // Large finite ranges can keep the front-end reservation alive for a long time. Check in + // batches so timeout handling does not add a clock read for every partition. + if (iterations++ % 1024 == 0) { + checkTimePartitionQueryTimeout(context); + } long curLeft = timeRangeList.get(index).getMin(); long curRight = timeRangeList.get(index).getMax(); if (curLeft >= endTime) { @@ -2354,13 +2363,26 @@ public static Pair, Pair> getTimePart return new Pair<>(result, new Pair<>(needLeftAll, needRightAll)); } + private static void checkTimePartitionQueryTimeout(MPPQueryContext context) { + long startTime = context.getStartTime(); + // Standalone analysis contexts may not have an execution deadline. + if (startTime == 0) { + return; + } + long currentTime = System.currentTimeMillis(); + long timeout = context.getTimeOut(); + if (currentTime - startTime >= timeout) { + throw new QueryTimeoutRuntimeException(startTime, currentTime, timeout); + } + } + private static void reserveMemoryForTimePartitionSlot( long maxTime, long minTime, MPPQueryContext context) { if (maxTime == Long.MAX_VALUE || minTime == Long.MIN_VALUE) { return; } long size = TimePartitionUtils.getEstimateTimePartitionSize(minTime, maxTime); - context.reserveMemoryForFrontEnd(estimateTimePartitionSlotMemory(size)); + context.reserveMemoryForTimePartitionSlots(estimateTimePartitionSlotMemory(size)); } private static List normalizeTimeRanges(List timeRanges) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java index 66fd1f38f9111..0e7f59a432bbb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java @@ -79,7 +79,20 @@ public long getFallbackBytesInTotalForTest() { @Override public void reserveMemoryCumulatively(final long size) { - bytesToBeReserved += size; + try { + bytesToBeReserved = Math.addExact(bytesToBeReserved, size); + } catch (ArithmeticException e) { + // A saturated partition-size estimate can be Long.MAX_VALUE while an earlier batch is + // still pending. Reject it without wrapping the counter or changing that pending batch. + throw new OperatorMemoryNotEnoughException( + String.format( + DataNodeQueryMessages + .EXCEPTION_QUERY_ARG_CANNOT_RESERVE_ARG_BYTES_PENDING_FRONT_END_MEMORY_RESERVATION_EXCEEDS_LONG_MAX_VALUE_BYTES_D38CD6D8, + queryId.getId(), + size), + Long.MAX_VALUE, + LOCAL_EXECUTION_PLANNER.getFreeMemoryForOperators()); + } if (bytesToBeReserved >= MEMORY_BATCH_THRESHOLD) { try { reserveMemoryImmediately(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java new file mode 100644 index 0000000000000..8a03e7d20cceb --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java @@ -0,0 +1,107 @@ +/* + * 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.protocol.client; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.AbstractConfigNodeClient; +import org.apache.iotdb.commons.client.property.ThriftClientProperty; +import org.apache.iotdb.confignode.rpc.thrift.IConfigNodeRPCService; +import org.apache.iotdb.confignode.rpc.thrift.TDataPartitionReq; +import org.apache.iotdb.confignode.rpc.thrift.TDataPartitionTableResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.transport.TTransportException; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ConfigNodeClientTest { + + @Test + public void testInvalidFrameDoesNotReconnectOrHideOriginalFailure() throws Exception { + IConfigNodeRPCService.Iface service = mock(IConfigNodeRPCService.Iface.class); + TDataPartitionReq request = new TDataPartitionReq(); + TTransportException failure = + new TTransportException(TTransportException.CORRUPTED_DATA, "Frame size exceeded"); + when(service.getDataPartitionTable(request)).thenThrow(failure); + AtomicInteger connections = new AtomicInteger(); + ConfigNodeClient client = newClient(service, connections, 1); + + assertSame( + failure, + assertThrows(TTransportException.class, () -> client.getDataPartitionTable(request))); + assertEquals(1, connections.get()); + verify(service).getDataPartitionTable(request); + } + + @Test + public void testTransientTransportFailureStillRetries() throws Exception { + IConfigNodeRPCService.Iface service = mock(IConfigNodeRPCService.Iface.class); + TDataPartitionReq request = new TDataPartitionReq(); + TDataPartitionTableResp response = + new TDataPartitionTableResp() + .setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + when(service.getDataPartitionTable(request)) + .thenThrow(new TTransportException(TTransportException.TIMED_OUT, "socket timeout")) + .thenReturn(response); + AtomicInteger connections = new AtomicInteger(); + ConfigNodeClient client = newClient(service, connections, 2); + + assertSame(response, client.getDataPartitionTable(request)); + assertEquals(2, connections.get()); + verify(service, times(2)).getDataPartitionTable(request); + } + + private ConfigNodeClient newClient( + IConfigNodeRPCService.Iface service, AtomicInteger connections, int allowedConnections) + throws Exception { + ConfigNodeClient client = + new ConfigNodeClient( + Collections.singletonList(new TEndPoint("127.0.0.1", 10710)), + new ThriftClientProperty.Builder().build(), + null) { + @Override + public void syncLatestConfigNodeList() { + // Keep the test candidates independent of the process-wide ConfigNodeInfo. + } + + @Override + public void connect(TEndPoint endpoint, int timeoutMs) { + assertTrue("Unexpected reconnect", connections.incrementAndGet() <= allowedConnections); + } + }; + Field field = AbstractConfigNodeClient.class.getDeclaredField("client"); + field.setAccessible(true); + field.set(client, service); + return client; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/common/MPPQueryContextTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/common/MPPQueryContextTest.java index fca6a3600cdf9..eca4ba5fed4a3 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/common/MPPQueryContextTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/common/MPPQueryContextTest.java @@ -22,6 +22,7 @@ import org.apache.iotdb.calc.exception.MemoryNotEnoughException; import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner; import org.apache.iotdb.db.queryengine.plan.planner.memory.OperatorMemoryNotEnoughException; +import org.apache.iotdb.db.utils.ErrorHandlingUtils; import org.apache.iotdb.rpc.TSStatusCode; import org.junit.Assert; @@ -32,6 +33,43 @@ public class MPPQueryContextTest { private static final long MEMORY_BATCH_THRESHOLD = 1024L * 1024L; + @Test + public void timePartitionMemoryFailureDoesNotUseColumnDiagnostics() { + MPPQueryContext context = new MPPQueryContext(new QueryId("time_partition_oom_test")); + context.initResultSetColumnMemoryTracking(10, 2, false); + context.recordMatchedSourceColumnsForResultSet(1); + context.recordGeneratedResultSetColumn(128); + long requestedBytes = requestLargerThanFreeOperatorMemory(); + MemoryNotEnoughException exception = + Assert.assertThrows( + MemoryNotEnoughException.class, + () -> context.reserveMemoryForTimePartitionSlots(requestedBytes)); + assertContains(exception.getMessage(), "while analyzing time partitions"); + assertContains(exception.getMessage(), "requested this time"); + assertContains(exception.getMessage(), "current free memory"); + Assert.assertFalse(exception.getMessage().contains("too many columns")); + Assert.assertFalse(exception.getMessage().contains("SLIMIT")); + Assert.assertEquals( + TSStatusCode.QUERY_EXECUTION_MEMORY_NOT_ENOUGH.getStatusCode(), + ErrorHandlingUtils.onQueryException(exception, "query").getCode()); + } + + @Test + public void laterPlanningMemoryFailureDoesNotUseStaleColumnStatistics() { + MPPQueryContext context = new MPPQueryContext(new QueryId("planning_oom_test")); + context.initResultSetColumnMemoryTracking(10, 2, false); + context.recordMatchedSourceColumnsForResultSet(1); + context.recordGeneratedResultSetColumn(128); + context.finishResultSetColumnMemoryTracking(); + long requestedBytes = requestLargerThanFreeOperatorMemory(); + MemoryNotEnoughException exception = + Assert.assertThrows( + MemoryNotEnoughException.class, () -> context.reserveMemoryForFrontEnd(requestedBytes)); + Assert.assertTrue(exception instanceof OperatorMemoryNotEnoughException); + Assert.assertFalse(exception.getMessage().contains("too many columns")); + Assert.assertFalse(exception.getMessage().contains("SLIMIT")); + } + @Test public void resultSetColumnMemoryNotEnoughExceptionContainsColumnDiagnostics() { MPPQueryContext context = new MPPQueryContext(new QueryId("result_column_oom_test")); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java new file mode 100644 index 0000000000000..292eab318e003 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java @@ -0,0 +1,300 @@ +/* + * 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.queryengine.plan.analyze; + +import org.apache.iotdb.calc.exception.MemoryNotEnoughException; +import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot; +import org.apache.iotdb.commons.exception.SemanticException; +import org.apache.iotdb.commons.partition.DataPartition; +import org.apache.iotdb.commons.partition.DataPartitionQueryParam; +import org.apache.iotdb.commons.path.PathPatternTree; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.utils.TimePartitionUtils; +import org.apache.iotdb.confignode.rpc.thrift.IConfigNodeRPCService; +import org.apache.iotdb.confignode.rpc.thrift.TDataPartitionReq; +import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; +import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.schematree.ClusterSchemaTree; +import org.apache.iotdb.db.queryengine.plan.Coordinator; +import org.apache.iotdb.db.queryengine.plan.parser.StatementGenerator; +import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.utils.ErrorHandlingUtils; +import org.apache.iotdb.rpc.TElasticFramedTransport; +import org.apache.iotdb.rpc.TSStatusCode; + +import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.Logger; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.read.ListAppender; +import org.apache.thrift.TException; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.transport.TMemoryBuffer; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.LoggerFactory; + +import java.time.ZoneOffset; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class QueryPartitionMemoryTest { + + private int originalDataNodeId; + private Logger logger; + private Level originalLogLevel; + private ListAppender appender; + + @Before + public void setUp() { + originalDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + // Exclude the coordinator's one-time startup event from per-query log assertions. + Coordinator.getInstance(); + logger = (Logger) LoggerFactory.getLogger(Coordinator.class); + originalLogLevel = logger.getLevel(); + logger.setLevel(Level.INFO); + appender = new ListAppender<>(); + appender.start(); + logger.addAppender(appender); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(originalDataNodeId); + logger.detachAppender(appender); + appender.stop(); + logger.setLevel(originalLogLevel); + } + + @Test + public void testPartitionMemoryFailureReportsItsSourceAndSql() { + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + String sql = "select s1 from root.sg.d1 where time >= 0 and time <= " + (Long.MAX_VALUE - 1); + for (long queryId = 200; queryId < 203; queryId++) { + final long id = queryId; + FakeSchemaFetcherImpl schemaFetcher = reservingSchemaFetcher(false); + MemoryNotEnoughException failure = + assertThrows(MemoryNotEnoughException.class, () -> execute(sql, id, schemaFetcher)); + assertTrue(failure.getMessage(), failure.getMessage().contains("time partitions")); + assertFalse(failure.getMessage(), failure.getMessage().contains("too many columns")); + assertFalse(failure.getMessage(), failure.getMessage().contains("SLIMIT")); + assertEquals( + TSStatusCode.QUERY_EXECUTION_MEMORY_NOT_ENOUGH.getStatusCode(), + ErrorHandlingUtils.onQueryException(failure, "query").getCode()); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(queryId)); + } + assertFailureLogs(sql, 3); + } + + @Test + public void testPartitionTimeoutLogsSqlAndReleasesEarlierReservations() { + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + String sql = "select s1 from root.sg.d1 where time >= 0 and time <= 100"; + FakeSchemaFetcherImpl schemaFetcher = reservingSchemaFetcher(true); + assertThrows(QueryTimeoutRuntimeException.class, () -> execute(sql, 300, schemaFetcher)); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(300L)); + assertFailureLogs(sql, 1); + } + + @Test + public void testSemanticFailureOnlyLogsSqlContext() { + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + String sql = "select count(s1), s1 from root.sg.d1"; + FakeSchemaFetcherImpl schemaFetcher = new FakeSchemaFetcherImpl(); + SemanticException failure = + assertThrows(SemanticException.class, () -> execute(sql, 600, schemaFetcher)); + assertEquals( + TSStatusCode.SEMANTIC_ERROR.getStatusCode(), + ErrorHandlingUtils.onQueryException(failure, "query").getCode()); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(600L)); + assertFailureLogs(sql, 1); + } + + @Test + public void testAnalysisErrorIsLoggedAndRethrownWithoutRetainingMemory() { + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + String sql = "select s1 from root.sg.d1 where time >= 0 and time <= 100"; + OutOfMemoryError failure = new OutOfMemoryError("simulated analysis allocation failure"); + FakeSchemaFetcherImpl schemaFetcher = + new FakeSchemaFetcherImpl() { + @Override + public ClusterSchemaTree fetchSchema( + PathPatternTree patternTree, + boolean withTemplate, + MPPQueryContext context, + boolean canSeeAuditDB) { + context.reserveMemoryForFrontEnd(2 * 1024 * 1024); + throw failure; + } + }; + assertSame( + failure, assertThrows(OutOfMemoryError.class, () -> execute(sql, 500, schemaFetcher))); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(500L)); + assertFailureLogs(sql, 1); + } + + @Test + public void testNormalQueryDoesNotLogFailure() { + String sql = "select s1 from root.sg.d1 where time > 10 and time < 5"; + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + try { + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + execute(sql, 400, new FakeSchemaFetcherImpl())); + assertTrue(appender.list.isEmpty()); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } finally { + Coordinator.getInstance().cleanupQueryExecution(400L); + } + } + + private int execute(String sql, long queryId, FakeSchemaFetcherImpl schemaFetcher) { + return Coordinator.getInstance() + .executeForTreeModel( + StatementGenerator.createStatement(sql, ZoneOffset.UTC), + queryId, + new SessionInfo(0, "root", ZoneOffset.UTC), + sql, + new FakePartitionFetcherImpl(), + schemaFetcher, + 60_000, + true, + false) + .status + .getCode(); + } + + private FakeSchemaFetcherImpl reservingSchemaFetcher(boolean expire) { + return new FakeSchemaFetcherImpl() { + @Override + public ClusterSchemaTree fetchSchema( + PathPatternTree patternTree, + boolean withTemplate, + MPPQueryContext context, + boolean canSeeAuditDB) { + context.reserveMemoryForFrontEnd(2 * 1024 * 1024); + if (expire) { + context.setStartTime(1); + context.setTimeOut(1); + } + return super.fetchSchema(patternTree, withTemplate, context, canSeeAuditDB); + } + }; + } + + private void assertFailureLogs(String sql, int count) { + assertEquals(count, appender.list.size()); + for (ILoggingEvent event : appender.list) { + assertEquals(Level.INFO, event.getLevel()); + assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains(sql)); + assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains("queryId:")); + assertNull(event.getThrowableProxy()); + } + } + + @Test + public void testOversizedPartitionRequestReleasesFrontEndMemory() { + LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); + long freeMemory = planner.getFreeMemoryForOperators(); + String sql = + "select s1 from root.sg.d1 where time >= 0 and time <= " + + TimePartitionUtils.getTimePartitionInterval() * 50_000; + for (long queryId = 100; queryId < 103; queryId++) { + final int[] serializedBytes = {0}; + IPartitionFetcher partitionFetcher = + new FakePartitionFetcherImpl() { + @Override + public DataPartition getDataPartition( + Map> params) { + assertTrue(planner.getFreeMemoryForOperators() < freeMemory); + DataPartitionQueryParam param = params.get("root.sg").get(0); + TTimeSlotList slots = + new TTimeSlotList(param.getTimePartitionSlotList(), false, false); + TDataPartitionReq request = + new TDataPartitionReq( + Collections.singletonMap( + "root.sg", Collections.singletonMap(new TSeriesPartitionSlot(0), slots))); + try (TElasticFramedTransport transport = + new TElasticFramedTransport(new TMemoryBuffer(0), 16, 1024, false) { + @Override + public void close() { + serializedBytes[0] = Math.max(serializedBytes[0], writeBuffer.getPos()); + super.close(); + } + }) { + new IConfigNodeRPCService.Client(new TBinaryProtocol(transport)) + .send_getDataPartitionTable(request); + fail("The partition request must exceed the frame size limit"); + } catch (TException e) { + throw new StatementAnalyzeException(e.getMessage()); + } + return null; + } + }; + final long id = queryId; + Coordinator coordinator = Coordinator.getInstance(); + Statement statement = StatementGenerator.createStatement(sql, ZoneOffset.UTC); + SessionInfo session = new SessionInfo(0, "root", ZoneOffset.UTC); + FakeSchemaFetcherImpl schemaFetcher = new FakeSchemaFetcherImpl(); + StatementAnalyzeException failure = + assertThrows( + StatementAnalyzeException.class, + () -> + coordinator.executeForTreeModel( + statement, + id, + session, + sql, + partitionFetcher, + schemaFetcher, + 60_000, + true, + false)); + assertTrue(failure.getMessage(), failure.getMessage().contains("at least")); + assertTrue(failure.getMessage(), failure.getMessage().contains("1024 bytes")); + assertTrue( + failure.getMessage(), failure.getMessage().contains("configured Thrift max frame size")); + assertFalse(failure.getMessage(), failure.getMessage().contains("dn_thrift_max_frame_size")); + assertTrue(serializedBytes[0] > 0); + assertTrue("Serialization must stop at the frame limit", serializedBytes[0] <= 1024); + assertEquals(freeMemory, planner.getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(queryId)); + } + assertFailureLogs(sql, 3); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryTimePartitionTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryTimePartitionTest.java index aa31e0ed3762d..bfd4556e09c90 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryTimePartitionTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryTimePartitionTest.java @@ -21,8 +21,10 @@ import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.utils.TimePartitionUtils; +import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner; import org.apache.tsfile.read.common.TimeRange; import org.apache.tsfile.read.filter.basic.Filter; @@ -38,10 +40,55 @@ import static org.apache.iotdb.db.queryengine.plan.analyze.AnalyzeVisitor.getTimePartitionSlotList; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; public class QueryTimePartitionTest { + @Test + public void testTimeoutDuringTimePartitionEnumerationReleasesMemory() { + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + MPPQueryContext context = + new MPPQueryContext(new QueryId("time_partition_enumeration_timeout")) { + private int timeoutChecks; + + @Override + public long getTimeOut() { + // Expire after the initial check, once the partition list has reserved memory. + return timeoutChecks++ == 0 ? Long.MAX_VALUE : 1; + } + }; + context.setStartTime(1); + Filter filter = + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000); + try { + assertThrows( + QueryTimeoutRuntimeException.class, () -> getTimePartitionSlotList(filter, context)); + assertTrue(LocalExecutionPlanner.getInstance().getFreeMemoryForOperators() < freeMemory); + } finally { + context.releaseAllMemoryReservedForFrontEnd(); + } + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } + + @Test + public void testExpiredQueryDoesNotEnumerateTimePartitions() { + MPPQueryContext context = new MPPQueryContext(new QueryId("expired_time_partitions")); + context.setStartTime(System.currentTimeMillis() - 10_000); + context.setTimeOut(1); + long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + Filter filter = + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000); + try { + assertThrows( + QueryTimeoutRuntimeException.class, () -> getTimePartitionSlotList(filter, context)); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } finally { + context.releaseAllMemoryReservedForFrontEnd(); + } + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } + @Test public void testAndTimeFilter() { Filter left = TimeFilterApi.gt(10); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java index 443d68e58830e..91f5772b58e26 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java @@ -40,6 +40,28 @@ public class LocalExecutionPlannerOperatorsMemoryTest { private long bytesHeldByTest = 0L; + @Test + public void testPendingReservationOverflowDoesNotCorruptAccounting() { + for (boolean highestPriority : new boolean[] {false, true}) { + long freeBefore = PLANNER.getFreeMemoryForOperators(); + NotThreadSafeMemoryReservationManager manager = + new NotThreadSafeMemoryReservationManager(new QueryId("overflow_query"), "test"); + manager.setHighestPriority(highestPriority); + try { + manager.reserveMemoryCumulatively(128); + Assert.assertThrows( + MemoryNotEnoughException.class, + () -> manager.reserveMemoryCumulatively(Long.MAX_VALUE)); + manager.reserveMemoryImmediately(); + Assert.assertEquals(128, manager.getReservedBytesInTotalForTest()); + Assert.assertEquals(freeBefore - 128, PLANNER.getFreeMemoryForOperators()); + } finally { + manager.releaseAllReservedMemory(); + } + Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); + } + } + @After public void tearDown() { if (bytesHeldByTest > 0) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/AbstractConfigNodeClient.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/AbstractConfigNodeClient.java index 2d20166fbb56c..fed868420d699 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/AbstractConfigNodeClient.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/AbstractConfigNodeClient.java @@ -192,6 +192,7 @@ import org.apache.thrift.TException; import org.apache.thrift.transport.TTransport; +import org.apache.thrift.transport.TTransportException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -439,6 +440,12 @@ protected final R executeRemoteCallWithRetry( } detectedNodeNum++; } catch (TException e) { + if (e instanceof TTransportException transportException + && transportException.getType() == TTransportException.CORRUPTED_DATA) { + // Reconnecting cannot repair an invalid/oversized frame. Repeated serialization would + // retain query memory and eventually hide the original error as a connection failure. + throw e; + } final String message = String.format( ConfigMessages.MSG_RECONNECTION_NODE_FAIL,