From bad9ce03b086cd9086b27df17c0a73760ca264da Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Sat, 14 Nov 2026 11:27:55 +0800 Subject: [PATCH 1/2] Fix query partition memory retention and failure diagnostics Port apache/iotdb#19269 from dev/1.3 to master. Preserve master RPC memory accounting, temporary large-frame buffers, and message-size tracking; apply non-retryable frame handling in the shared ConfigNode client. Adapt regression tests to current APIs and error codes, and provide both English and Chinese diagnostics. (cherry picked from commit 33277a6cf5c064d31810c519f5e1f7d4c4c1e24b) (cherry picked from commit 6a828ef41a41879634862e6eb0da178d0b37c28b) (cherry picked from commit 27c44f50a2fafc36000910682a7467c7ca0498ae) (cherry picked from commit 856d4a0d820e0aa115f5ffdeae4b44687c26fff9) (cherry picked from commit 827763a31d9de709c072d24097707738dd189c47) (cherry picked from commit 86b83e094fa9797bb9503bf108089d371a63260c) (cherry picked from commit c7003fff631e252ae4f0796ae258124caffcf2a5) --- .github/workflows/sonar-codecov.yml | 5 + .../apache/iotdb/rpc/i18n/RpcMessages.java | 15 + .../apache/iotdb/rpc/i18n/RpcMessages.java | 15 + .../apache/iotdb/rpc/AutoResizingBuffer.java | 29 +- .../rpc/AutoScalingBufferWriteTransport.java | 49 ++- .../TCompressedElasticFramedTransport.java | 22 ++ .../iotdb/rpc/TElasticFramedTransport.java | 35 +- .../iotdb/rpc/AutoResizingBufferTest.java | 101 ++++++ .../rpc/TElasticFramedTransportTest.java | 210 ++++++++++++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 9 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 9 + .../queryengine/common/MPPQueryContext.java | 29 ++ .../db/queryengine/plan/Coordinator.java | 9 + .../plan/analyze/AnalyzeVisitor.java | 24 +- ...NotThreadSafeMemoryReservationManager.java | 15 +- .../protocol/client/ConfigNodeClientTest.java | 107 +++++++ .../common/MPPQueryContextTest.java | 38 +++ .../analyze/QueryPartitionMemoryTest.java | 300 ++++++++++++++++++ .../plan/analyze/QueryTimePartitionTest.java | 47 +++ ...alExecutionPlannerOperatorsMemoryTest.java | 22 ++ .../client/AbstractConfigNodeClient.java | 7 + 21 files changed, 1081 insertions(+), 16 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java diff --git a/.github/workflows/sonar-codecov.yml b/.github/workflows/sonar-codecov.yml index 88f69e5938d4f..eab04b2252b4d 100644 --- a/.github/workflows/sonar-codecov.yml +++ b/.github/workflows/sonar-codecov.yml @@ -50,6 +50,11 @@ jobs: steps: - uses: actions/checkout@v5 + - name: Set up JDK 21 + uses: actions/setup-java@v5 + with: + distribution: corretto + java-version: 21 - name: Cache Maven packages uses: actions/cache@v5 with: 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..b9197cb7d27fd 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,19 @@ 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."; + } 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..7cb7c584e1cf4 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,19 @@ 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 字节。"; + } 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..d9eebc7fc0d96 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,29 @@ 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. + close(); + String message = + 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; + throw new TTransportException( + TTransportException.CORRUPTED_DATA, String.format(message, size, maxSize)); + } + 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..9ad6152132697 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,107 @@ 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 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..a2b77aa8c9cff 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 @@ -25,16 +25,226 @@ import org.apache.thrift.transport.TTransportException; import org.junit.Test; +import java.lang.reflect.Field; 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 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, From 6185c343266fda18ed1f06710553d8463d41521c Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Sat, 14 Nov 2026 15:22:34 +0800 Subject: [PATCH 2/2] Preserve RPC rejection endpoints and cover reuse threshold Capture the destination before closing a socket when rejecting oversized writes, and provide complete English and Chinese diagnostics. Cover writes that cross the reuse threshold while fitting the current array, and remove the coverage JDK workaround that master does not need. --- .github/workflows/sonar-codecov.yml | 5 --- .../apache/iotdb/rpc/i18n/RpcMessages.java | 6 +++ .../apache/iotdb/rpc/i18n/RpcMessages.java | 6 +++ .../iotdb/rpc/TElasticFramedTransport.java | 37 ++++++++++++---- .../iotdb/rpc/AutoResizingBufferTest.java | 30 +++++++++++++ .../rpc/TElasticFramedTransportTest.java | 42 +++++++++++++++++++ 6 files changed, 113 insertions(+), 13 deletions(-) diff --git a/.github/workflows/sonar-codecov.yml b/.github/workflows/sonar-codecov.yml index eab04b2252b4d..88f69e5938d4f 100644 --- a/.github/workflows/sonar-codecov.yml +++ b/.github/workflows/sonar-codecov.yml @@ -50,11 +50,6 @@ jobs: steps: - uses: actions/checkout@v5 - - name: Set up JDK 21 - uses: actions/setup-java@v5 - with: - distribution: corretto - java-version: 21 - name: Cache Maven packages uses: actions/cache@v5 with: 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 b9197cb7d27fd..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 @@ -142,4 +142,10 @@ private RpcMessages() {} 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 7cb7c584e1cf4..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 @@ -123,4 +123,10 @@ private RpcMessages() {} 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/TElasticFramedTransport.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TElasticFramedTransport.java index d9eebc7fc0d96..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 @@ -394,15 +394,36 @@ public void write(byte[] buf, int off, int len) throws TTransportException { 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 = - 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; - throw new TTransportException( - TTransportException.CORRUPTED_DATA, String.format(message, size, maxSize)); + 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); } 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 9ad6152132697..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 @@ -112,6 +112,36 @@ public void testFramedFastPathReservesMemoryAfterReopening() throws Exception { 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); 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 a2b77aa8c9cff..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,10 +22,14 @@ 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; @@ -41,6 +45,44 @@ 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}) {