From 33277a6cf5c064d31810c519f5e1f7d4c4c1e24b Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 21:23:19 +0800 Subject: [PATCH 1/7] Fix query memory retention during oversized partition RPC serialization --- .../apache/iotdb/rpc/AutoResizingBuffer.java | 17 ++- .../rpc/AutoScalingBufferWriteTransport.java | 14 +- .../TCompressedElasticFramedTransport.java | 7 + .../iotdb/rpc/TElasticFramedTransport.java | 21 ++- .../iotdb/rpc/AutoResizingBufferTest.java | 42 ++++++ .../rpc/TElasticFramedTransportTest.java | 60 ++++++++ .../plan/analyze/AnalyzeVisitor.java | 21 +++ .../analyze/QueryPartitionMemoryTest.java | 129 ++++++++++++++++++ .../plan/analyze/QueryTimePartitionTest.java | 49 +++++++ 9 files changed, 350 insertions(+), 10 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java 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 ca47d80acae23..d933ecb1b40a5 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 @@ -34,6 +34,8 @@ */ class AutoResizingBuffer { + 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; @@ -59,16 +61,23 @@ public void resizeIfNecessary(int size) { } void ensureCapacity(int size) { + if (size < 0 || size > MAX_ARRAY_SIZE) { + throw new IllegalArgumentException("Invalid buffer size: " + size); + } 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); - this.array = Arrays.copyOf(array, newCapacity); + this.array = Arrays.copyOf(array, calculateNewCapacity(currentCapacity, 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; } 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 9e57e53768f95..ccb70b6c6024a 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 @@ -43,10 +43,20 @@ public int read(byte[] buf, int off, int len) { } @Override - public void write(byte[] toWrite, int off, int len) { + public void write(byte[] toWrite, int off, int len) throws TTransportException { + long requiredCapacity = (long) pos + len; + if (requiredCapacity > AutoResizingBuffer.MAX_ARRAY_SIZE) { + throw new TTransportException( + TTransportException.CORRUPTED_DATA, + "Write buffer size (" + + requiredCapacity + + ") larger than maximum array size (" + + AutoResizingBuffer.MAX_ARRAY_SIZE + + ")!"); + } // A small prefix during serialization does not indicate a small frame. Only the complete // frame size, supplied after flushing, may trigger shrinking. - buf.ensureCapacity(pos + len); + buf.ensureCapacity((int) requiredCapacity); System.arraycopy(toWrite, off, buf.array(), pos, len); pos += len; } 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 ab00ca2ac6393..2c0e74b37d0fb 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 @@ -40,6 +40,13 @@ protected TCompressedElasticFramedTransport( readCompressBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize); } + @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 { underlying.readAll(i32buf, 0, 4); 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 cbc74e22f42c6..f8a6f5af58eef 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 @@ -40,9 +40,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; @@ -288,10 +287,24 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException { } @Override - public void write(byte[] buf, int off, int len) { + public void write(byte[] buf, int off, int len) throws TTransportException { + 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(); + throw new TTransportException( + TTransportException.CORRUPTED_DATA, + "Frame size (" + size + ") larger than protect max size (" + maxSize + ")!"); + } writeBuffer.write(buf, off, len); } + protected int getMaxWriteBufferSize() { + return Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE); + } + public TTransport getSocket() { return underlying; } 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 d4ddaa003138d..4933a0e0e9662 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 @@ -20,6 +20,7 @@ package org.apache.iotdb.rpc; import org.apache.thrift.transport.TMemoryBuffer; +import org.apache.thrift.transport.TTransportException; import org.junit.Assert; import org.junit.Test; @@ -28,6 +29,47 @@ public class AutoResizingBufferTest { + @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() { + 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() { AutoResizingBuffer buffer = new AutoResizingBuffer(100); 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 086dc33825065..55b3f1418cb64 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 @@ -20,17 +20,77 @@ package org.apache.iotdb.rpc; import org.apache.thrift.transport.TByteBuffer; +import org.apache.thrift.transport.TMemoryBuffer; import org.apache.thrift.transport.TTransportException; import org.junit.Test; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +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().contains("Frame size (65)")); + assertSame(buffer, transport.writeBuffer.getBuffer()); + assertEquals(prefixLength, transport.writeBuffer.getPos()); + assertEquals(0, underlying.length()); + assertTrue(closed[0]); + } + } + } + + @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), 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); + } + @Test public void testSingularSize() { 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 c393c275dd529..b9ab605f5a72d 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 @@ -47,6 +47,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.SemanticException; import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; @@ -2327,6 +2328,7 @@ private DataPartition fetchDataPartitionByDevices( @SuppressWarnings("squid:S3776") // Suppress high Cognitive Complexity warning public static Pair, Pair> getTimePartitionSlotList( Filter timeFilter, MPPQueryContext context) { + checkTimePartitionQueryTimeout(context); if (timeFilter == null) { // (-oo, +oo) return new Pair<>(Collections.emptyList(), new Pair<>(true, true)); @@ -2370,7 +2372,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) { @@ -2405,6 +2413,19 @@ 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) { 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..78dec77e80473 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/QueryPartitionMemoryTest.java @@ -0,0 +1,129 @@ +/* + * 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.common.rpc.thrift.TSeriesPartitionSlot; +import org.apache.iotdb.commons.partition.DataPartition; +import org.apache.iotdb.commons.partition.DataPartitionQueryParam; +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.sql.StatementAnalyzeException; +import org.apache.iotdb.db.queryengine.common.SessionInfo; +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.rpc.TElasticFramedTransport; + +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 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.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class QueryPartitionMemoryTest { + + private int originalDataNodeId; + + @Before + public void setUp() { + originalDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(originalDataNodeId); + } + + @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] = 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; + } + }; + try { + Coordinator.getInstance() + .executeForTreeModel( + StatementGenerator.createStatement(sql, ZoneOffset.UTC), + queryId, + new SessionInfo(0, "root", ZoneOffset.UTC, ""), + sql, + partitionFetcher, + new FakeSchemaFetcherImpl(), + 60_000, + true, + false); + fail("The oversized partition request must fail query analysis"); + } catch (StatementAnalyzeException e) { + assertTrue(e.getMessage(), e.getMessage().contains("larger than protect max size (1024)")); + } + assertTrue(serializedBytes[0] > 0); + assertTrue("Serialization must stop at the frame limit", serializedBytes[0] <= 1024); + assertEquals(freeMemory, planner.getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(queryId)); + } + } +} 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..5293445498be1 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; @@ -39,9 +41,56 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; 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); + try { + getTimePartitionSlotList( + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000), + context); + fail("Time partition enumeration must periodically check the query timeout"); + } catch (QueryTimeoutRuntimeException expected) { + 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(); + try { + getTimePartitionSlotList( + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000), + context); + fail("An expired query must stop before enumerating time partitions"); + } catch (QueryTimeoutRuntimeException expected) { + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } finally { + context.releaseAllMemoryReservedForFrontEnd(); + } + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + } + @Test public void testAndTimeFilter() { Filter left = TimeFilterApi.gt(10); From 6a828ef41a41879634862e6eb0da178d0b37c28b Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 21:48:22 +0800 Subject: [PATCH 2/7] Fix query partition failure diagnostics and log failing SQL --- .../TCompressedElasticFramedTransport.java | 12 ++ .../rpc/TElasticFramedTransportTest.java | 42 +++++ .../db/protocol/client/ConfigNodeClient.java | 6 + .../queryengine/common/MPPQueryContext.java | 31 ++++ .../db/queryengine/plan/Coordinator.java | 5 + .../plan/analyze/AnalyzeVisitor.java | 3 +- ...NotThreadSafeMemoryReservationManager.java | 15 +- .../protocol/client/ConfigNodeClientTest.java | 101 +++++++++++ .../common/MPPQueryContextTest.java | 37 +++++ .../analyze/QueryPartitionMemoryTest.java | 157 ++++++++++++++++-- .../plan/analyze/QueryTimePartitionTest.java | 20 +-- ...alExecutionPlannerOperatorsMemoryTest.java | 22 +++ 12 files changed, 422 insertions(+), 29 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java 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 2c0e74b37d0fb..ff2037d81afda 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 @@ -56,6 +56,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); @@ -73,6 +74,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); @@ -90,6 +92,16 @@ 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, "Invalid compression buffer size: " + 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/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java b/iotdb-client/service-rpc/src/test/java/org/apache/iotdb/rpc/TElasticFramedTransportTest.java index 55b3f1418cb64..3cdfcc334256e 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 @@ -30,6 +30,7 @@ import static org.junit.Assert.assertArrayEquals; 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.junit.Assert.fail; @@ -91,6 +92,47 @@ public void testCompressedFrameLimitAppliesAfterCompression() throws Exception { assertArrayEquals(payload, actual); } + @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 testSingularSize() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java index cf83a81cd4890..303323cfaff99 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java @@ -386,6 +386,12 @@ private T executeRemoteCallWithRetry(final Operation call, final Predicat } detectedNodeNum++; } catch (TException e) { + if (e instanceof TTransportException + && ((TTransportException) e).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( MSG_RECONNECTION_DATANODE_FAIL, 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 73d4093e6c7d6..dd16d0d3a9060 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 @@ -460,6 +460,37 @@ public void reserveMemoryForFrontEndImmediately() { } } + public void reserveMemoryForTimePartitionSlots(final long bytes) { + try { + this.memoryReservationManager.reserveMemoryCumulatively(bytes); + } catch (MemoryNotEnoughException e) { + long requestedBytes = + e instanceof OperatorMemoryNotEnoughException + ? ((OperatorMemoryNotEnoughException) e).getRequestedBytes() + : bytes; + long freeBytes = + e instanceof OperatorMemoryNotEnoughException + ? ((OperatorMemoryNotEnoughException) e).getFreeBytes() + : LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + throw new MemoryNotEnoughException( + String.format( + Locale.ROOT, + "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", + 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 22951f71258e0..dc1f8106c82aa 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 @@ -162,6 +162,11 @@ private ExecutionResult execution( result.status.setNeedRetry(true); } return result; + } catch (RuntimeException | Error e) { + // Analysis can fail before the execution is registered, so later error handling cannot look + // up its SQL in queryExecutionMap. Keep the SQL and query ID in the failure log here. + LOGGER.warn("[QueryFailed] queryId: {}, sql: {}", globalQueryId, sql, e); + 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 b9ab605f5a72d..cf7ccc884adf4 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 @@ -2329,6 +2329,7 @@ private DataPartition fetchDataPartitionByDevices( 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)); @@ -2432,7 +2433,7 @@ private static void reserveMemoryForTimePartitionSlot( 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 514a2935f6ce3..1140fec0793e6 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 @@ -77,7 +77,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( + "Query " + + queryId.getId() + + " cannot reserve " + + size + + " bytes: pending front-end memory reservation exceeds Long.MAX_VALUE bytes.", + 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..edef1e7202fd5 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeClientTest.java @@ -0,0 +1,101 @@ +/* + * 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.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 connect(TEndPoint endpoint, int timeoutMs) { + assertTrue("Unexpected reconnect", connections.incrementAndGet() <= allowedConnections); + } + }; + Field field = ConfigNodeClient.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 6d5cd68e04989..711e98dfff4cb 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 @@ -33,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.QUOTA_MEM_QUERY_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 index 78dec77e80473..8644991c24e56 100644 --- 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 @@ -16,29 +16,43 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.iotdb.db.queryengine.plan.analyze; import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot; 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.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.SessionInfo; +import org.apache.iotdb.db.queryengine.common.schematree.ClusterSchemaTree; +import org.apache.iotdb.db.queryengine.exception.MemoryNotEnoughException; 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; @@ -46,23 +60,127 @@ import java.util.Map; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +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); + logger = (Logger) LoggerFactory.getLogger(Coordinator.class); + originalLogLevel = logger.getLevel(); + logger.setLevel(Level.WARN); + 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.QUOTA_MEM_QUERY_NOT_ENOUGH.getStatusCode(), + ErrorHandlingUtils.onQueryException(failure, "query").getCode()); + assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); + assertNull(Coordinator.getInstance().getQueryExecution(queryId)); + } + assertFailureLogs(sql, MemoryNotEnoughException.class, 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, QueryTimeoutRuntimeException.class, 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) { + context.reserveMemoryForFrontEnd(2 * 1024 * 1024); + if (expire) { + context.setStartTime(1); + context.setTimeOut(1); + } + return super.fetchSchema(patternTree, withTemplate, context); + } + }; + } + + private void assertFailureLogs(String sql, Class exception, int count) { + assertEquals(count, appender.list.size()); + for (ILoggingEvent event : appender.list) { + assertEquals(Level.WARN, event.getLevel()); + assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains(sql)); + assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains("queryId:")); + assertNotNull(event.getThrowableProxy()); + assertEquals(exception.getName(), event.getThrowableProxy().getClassName()); + } } @Test @@ -104,26 +222,33 @@ public void close() { return null; } }; - try { - Coordinator.getInstance() - .executeForTreeModel( - StatementGenerator.createStatement(sql, ZoneOffset.UTC), - queryId, - new SessionInfo(0, "root", ZoneOffset.UTC, ""), - sql, - partitionFetcher, - new FakeSchemaFetcherImpl(), - 60_000, - true, - false); - fail("The oversized partition request must fail query analysis"); - } catch (StatementAnalyzeException e) { - assertTrue(e.getMessage(), e.getMessage().contains("larger than protect max size (1024)")); - } + 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("larger than protect max size (1024)")); 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, StatementAnalyzeException.class, 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 5293445498be1..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 @@ -40,8 +40,8 @@ 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; -import static org.junit.Assert.fail; public class QueryTimePartitionTest { @@ -59,12 +59,11 @@ public long getTimeOut() { } }; context.setStartTime(1); + Filter filter = + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000); try { - getTimePartitionSlotList( - TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000), - context); - fail("Time partition enumeration must periodically check the query timeout"); - } catch (QueryTimeoutRuntimeException expected) { + assertThrows( + QueryTimeoutRuntimeException.class, () -> getTimePartitionSlotList(filter, context)); assertTrue(LocalExecutionPlanner.getInstance().getFreeMemoryForOperators() < freeMemory); } finally { context.releaseAllMemoryReservedForFrontEnd(); @@ -78,12 +77,11 @@ public void testExpiredQueryDoesNotEnumerateTimePartitions() { context.setStartTime(System.currentTimeMillis() - 10_000); context.setTimeOut(1); long freeMemory = LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + Filter filter = + TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000); try { - getTimePartitionSlotList( - TimeFilterApi.between(0, TimePartitionUtils.getTimePartitionInterval() * 50_000), - context); - fail("An expired query must stop before enumerating time partitions"); - } catch (QueryTimeoutRuntimeException expected) { + assertThrows( + QueryTimeoutRuntimeException.class, () -> getTimePartitionSlotList(filter, context)); assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); } finally { context.releaseAllMemoryReservedForFrontEnd(); 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 cc763763d51e8..97264e25260bc 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) { From 27c44f50a2fafc36000910682a7467c7ca0498ae Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 21:53:48 +0800 Subject: [PATCH 3/7] Pin coverage workflow to JDK 21 for formatter compatibility --- .github/workflows/sonar-codecov.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/sonar-codecov.yml b/.github/workflows/sonar-codecov.yml index e29ff35cf67fe..e633a5fbcdc88 100644 --- a/.github/workflows/sonar-codecov.yml +++ b/.github/workflows/sonar-codecov.yml @@ -43,6 +43,11 @@ jobs: steps: - uses: actions/checkout@v4 + - name: Set up JDK 21 + uses: actions/setup-java@v4 + with: + distribution: corretto + java-version: 21 - name: Cache Maven packages uses: actions/cache@v4 with: From 856d4a0d820e0aa115f5ffdeae4b44687c26fff9 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 22:05:58 +0800 Subject: [PATCH 4/7] Address static analysis feedback and verify error cleanup --- .../db/protocol/client/ConfigNodeClient.java | 4 ++-- .../queryengine/common/MPPQueryContext.java | 16 +++++++------- .../db/queryengine/plan/Coordinator.java | 3 +++ .../analyze/QueryPartitionMemoryTest.java | 22 +++++++++++++++++++ 4 files changed, 35 insertions(+), 10 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java index 303323cfaff99..07d12e5b60733 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java @@ -386,8 +386,8 @@ private T executeRemoteCallWithRetry(final Operation call, final Predicat } detectedNodeNum++; } catch (TException e) { - if (e instanceof TTransportException - && ((TTransportException) e).getType() == TTransportException.CORRUPTED_DATA) { + 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; 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 dd16d0d3a9060..7db3b7abe4636 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 @@ -464,14 +464,14 @@ public void reserveMemoryForTimePartitionSlots(final long bytes) { try { this.memoryReservationManager.reserveMemoryCumulatively(bytes); } catch (MemoryNotEnoughException e) { - long requestedBytes = - e instanceof OperatorMemoryNotEnoughException - ? ((OperatorMemoryNotEnoughException) e).getRequestedBytes() - : bytes; - long freeBytes = - e instanceof OperatorMemoryNotEnoughException - ? ((OperatorMemoryNotEnoughException) e).getFreeBytes() - : LocalExecutionPlanner.getInstance().getFreeMemoryForOperators(); + 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, 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 dc1f8106c82aa..0b8615994d5bd 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 @@ -124,6 +124,9 @@ private Coordinator() { new ThreadPoolExecutor.CallerRunsPolicy()); } + // This boundary still owns the SQL. Observe even Errors for diagnostics, then rethrow unchanged + // rather than recover or wrap them, preserving caller error handling and finally cleanup. + @SuppressWarnings({"java:S1181", "java:S2139"}) private ExecutionResult execution( long queryId, SessionInfo session, 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 index 8644991c24e56..3eb8acfa48e48 100644 --- 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 @@ -63,6 +63,7 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; 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; @@ -126,6 +127,27 @@ public void testPartitionTimeoutLogsSqlAndReleasesEarlierReservations() { assertFailureLogs(sql, QueryTimeoutRuntimeException.class, 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) { + 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, OutOfMemoryError.class, 1); + } + @Test public void testNormalQueryDoesNotLogFailure() { String sql = "select s1 from root.sg.d1 where time > 10 and time < 5"; From 827763a31d9de709c072d24097707738dd189c47 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 13 Nov 2026 23:17:37 +0800 Subject: [PATCH 5/7] Bound RPC write buffers and avoid duplicate failure stacks --- .../apache/iotdb/rpc/AutoResizingBuffer.java | 10 +++- .../rpc/AutoScalingBufferWriteTransport.java | 13 +++-- .../iotdb/rpc/TElasticFramedTransport.java | 20 ++++++-- .../rpc/TElasticFramedTransportTest.java | 47 ++++++++++++++++++- .../db/queryengine/plan/Coordinator.java | 12 ++--- .../analyze/QueryPartitionMemoryTest.java | 40 +++++++++++----- 6 files changed, 112 insertions(+), 30 deletions(-) diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoResizingBuffer.java index d933ecb1b40a5..ee86294cced09 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 @@ -61,12 +61,18 @@ public void resizeIfNecessary(int size) { } void ensureCapacity(int size) { - if (size < 0 || size > MAX_ARRAY_SIZE) { + ensureCapacity(size, MAX_ARRAY_SIZE); + } + + void ensureCapacity(int size, int maxCapacity) { + int limit = Math.min(maxCapacity, MAX_ARRAY_SIZE); + if (size < 0 || size > limit) { throw new IllegalArgumentException("Invalid buffer size: " + size); } final int currentCapacity = this.array.length; if (currentCapacity < size) { - this.array = Arrays.copyOf(array, calculateNewCapacity(currentCapacity, size)); + this.array = + Arrays.copyOf(array, Math.min(calculateNewCapacity(currentCapacity, size), limit)); bufTooLargeCounter = RpcUtils.MAX_BUFFER_OVERSIZE_TIME; } } 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 ccb70b6c6024a..d9501a0dd0c2c 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 @@ -44,19 +44,24 @@ 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 > AutoResizingBuffer.MAX_ARRAY_SIZE) { + if (requiredCapacity > limit) { throw new TTransportException( TTransportException.CORRUPTED_DATA, "Write buffer size (" + requiredCapacity - + ") larger than maximum array size (" - + AutoResizingBuffer.MAX_ARRAY_SIZE + + ") larger than maximum buffer size (" + + limit + ")!"); } // A small prefix during serialization does not indicate a small frame. Only the complete // frame size, supplied after flushing, may trigger shrinking. - buf.ensureCapacity((int) requiredCapacity); + buf.ensureCapacity((int) requiredCapacity, limit); System.arraycopy(toWrite, off, buf.array(), pos, len); pos += 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 f8a6f5af58eef..b469af38ff82a 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 @@ -87,7 +87,11 @@ public TElasticFramedTransport( this.thriftMaxFrameSize = thriftMaxFrameSize; this.copyBinary = copyBinary; readBuffer = new AutoScalingBufferReadTransport(thriftDefaultBufferSize); - writeBuffer = new AutoScalingBufferWriteTransport(thriftDefaultBufferSize); + writeBuffer = + new AutoScalingBufferWriteTransport( + Math.min( + thriftDefaultBufferSize, + Math.min(thriftMaxFrameSize, AutoResizingBuffer.MAX_ARRAY_SIZE))); } protected final int thriftDefaultBufferSize; @@ -294,11 +298,21 @@ public void write(byte[] buf, int off, int len) throws TTransportException { // 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 limitDescription = + maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE + ? "maximum Java array size" + : "RPC frame-size limit (dn_thrift_max_frame_size on DataNode)"; throw new TTransportException( TTransportException.CORRUPTED_DATA, - "Frame size (" + size + ") larger than protect max size (" + maxSize + ")!"); + "Serialized frame size is at least " + + size + + " bytes, exceeding the " + + limitDescription + + " of " + + maxSize + + " bytes."); } - writeBuffer.write(buf, off, len); + writeBuffer.write(buf, off, len, maxSize); } protected int getMaxWriteBufferSize() { 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 3cdfcc334256e..c15db63593176 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 @@ -26,6 +26,7 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.util.Arrays; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -56,7 +57,8 @@ public void close() { fail("An oversized frame must fail during serialization, before flush"); } catch (TTransportException e) { assertEquals(TTransportException.CORRUPTED_DATA, e.getType()); - assertTrue(e.getMessage().contains("Frame size (65)")); + assertTrue(e.getMessage(), e.getMessage().contains("at least 65 bytes")); + assertTrue(e.getMessage(), e.getMessage().contains("dn_thrift_max_frame_size")); assertSame(buffer, transport.writeBuffer.getBuffer()); assertEquals(prefixLength, transport.writeBuffer.getPos()); assertEquals(0, underlying.length()); @@ -65,6 +67,47 @@ public void close() { } } + @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 = @@ -83,7 +126,7 @@ public void testFrameAtSizeLimit() throws Exception { @Test public void testCompressedFrameLimitAppliesAfterCompression() throws Exception { TElasticFramedTransport transport = - new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 8, 64, false); + new TSnappyElasticFramedTransport(new TMemoryBuffer(0), 128, 64, false); byte[] payload = new byte[256]; transport.write(payload); transport.flush(); 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 0b8615994d5bd..00fd3344119c0 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 @@ -124,9 +124,8 @@ private Coordinator() { new ThreadPoolExecutor.CallerRunsPolicy()); } - // This boundary still owns the SQL. Observe even Errors for diagnostics, then rethrow unchanged - // rather than recover or wrap them, preserving caller error handling and finally cleanup. - @SuppressWarnings({"java:S1181", "java:S2139"}) + // Record the SQL for Errors as well, preserving the original failure and finally cleanup. + @SuppressWarnings("java:S1181") private ExecutionResult execution( long queryId, SessionInfo session, @@ -166,9 +165,10 @@ private ExecutionResult execution( } return result; } catch (RuntimeException | Error e) { - // Analysis can fail before the execution is registered, so later error handling cannot look - // up its SQL in queryExecutionMap. Keep the SQL and query ID in the failure log here. - LOGGER.warn("[QueryFailed] queryId: {}, sql: {}", globalQueryId, sql, 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("[QueryFailed] queryId: {}, sql: {}", globalQueryId, sql); throw e; } finally { if (queryContext != null) { 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 index 3eb8acfa48e48..071dc4fa6f57a 100644 --- 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 @@ -29,6 +29,7 @@ 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.SemanticException; import org.apache.iotdb.db.exception.sql.StatementAnalyzeException; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.SessionInfo; @@ -61,7 +62,6 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; @@ -81,7 +81,7 @@ public void setUp() { IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); logger = (Logger) LoggerFactory.getLogger(Coordinator.class); originalLogLevel = logger.getLevel(); - logger.setLevel(Level.WARN); + logger.setLevel(Level.INFO); appender = new ListAppender<>(); appender.start(); logger.addAppender(appender); @@ -113,7 +113,7 @@ public void testPartitionMemoryFailureReportsItsSourceAndSql() { assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); assertNull(Coordinator.getInstance().getQueryExecution(queryId)); } - assertFailureLogs(sql, MemoryNotEnoughException.class, 3); + assertFailureLogs(sql, 3); } @Test @@ -124,7 +124,22 @@ public void testPartitionTimeoutLogsSqlAndReleasesEarlierReservations() { assertThrows(QueryTimeoutRuntimeException.class, () -> execute(sql, 300, schemaFetcher)); assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); assertNull(Coordinator.getInstance().getQueryExecution(300L)); - assertFailureLogs(sql, QueryTimeoutRuntimeException.class, 1); + 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 @@ -145,7 +160,7 @@ public ClusterSchemaTree fetchSchema( failure, assertThrows(OutOfMemoryError.class, () -> execute(sql, 500, schemaFetcher))); assertEquals(freeMemory, LocalExecutionPlanner.getInstance().getFreeMemoryForOperators()); assertNull(Coordinator.getInstance().getQueryExecution(500L)); - assertFailureLogs(sql, OutOfMemoryError.class, 1); + assertFailureLogs(sql, 1); } @Test @@ -194,14 +209,13 @@ public ClusterSchemaTree fetchSchema( }; } - private void assertFailureLogs(String sql, Class exception, int count) { + private void assertFailureLogs(String sql, int count) { assertEquals(count, appender.list.size()); for (ILoggingEvent event : appender.list) { - assertEquals(Level.WARN, event.getLevel()); + assertEquals(Level.INFO, event.getLevel()); assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains(sql)); assertTrue(event.getFormattedMessage(), event.getFormattedMessage().contains("queryId:")); - assertNotNull(event.getThrowableProxy()); - assertEquals(exception.getName(), event.getThrowableProxy().getClassName()); + assertNull(event.getThrowableProxy()); } } @@ -263,14 +277,14 @@ public void close() { 60_000, true, false)); - assertTrue( - failure.getMessage(), - failure.getMessage().contains("larger than protect max size (1024)")); + assertTrue(failure.getMessage(), failure.getMessage().contains("at least")); + assertTrue(failure.getMessage(), failure.getMessage().contains("1024 bytes")); + assertTrue(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, StatementAnalyzeException.class, 3); + assertFailureLogs(sql, 3); } } From 86b83e094fa9797bb9503bf108089d371a63260c Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Sat, 14 Nov 2026 09:38:39 +0800 Subject: [PATCH 6/7] Keep RPC limit diagnostics generic and optimize buffered writes --- .../rpc/AutoScalingBufferWriteTransport.java | 12 ++++ .../iotdb/rpc/TElasticFramedTransport.java | 5 +- .../rpc/TElasticFramedTransportTest.java | 69 ++++++++++++++++++- .../analyze/QueryPartitionMemoryTest.java | 4 +- 4 files changed, 87 insertions(+), 3 deletions(-) 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 d9501a0dd0c2c..45dfc1083968d 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java @@ -47,6 +47,18 @@ public void write(byte[] toWrite, int off, int len) throws TTransportException { write(toWrite, off, len, AutoResizingBuffer.MAX_ARRAY_SIZE); } + // 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 = buf.array(); + if (len > array.length - pos) { + return false; + } + System.arraycopy(toWrite, off, array, pos, len); + pos += len; + return true; + } + 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; 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 b469af38ff82a..beff263faee8d 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 @@ -292,6 +292,9 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException { @Override public void write(byte[] buf, int off, int len) throws TTransportException { + if (writeBuffer.tryWriteWithinCapacity(buf, off, len)) { + return; + } long size = (long) writeBuffer.getPos() + len; int maxSize = getMaxWriteBufferSize(); if (size > maxSize) { @@ -301,7 +304,7 @@ public void write(byte[] buf, int off, int len) throws TTransportException { String limitDescription = maxSize == AutoResizingBuffer.MAX_ARRAY_SIZE ? "maximum Java array size" - : "RPC frame-size limit (dn_thrift_max_frame_size on DataNode)"; + : "configured Thrift max frame size"; throw new TTransportException( TTransportException.CORRUPTED_DATA, "Serialized frame size is at least " 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 c15db63593176..87b81ecbceb4b 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 @@ -24,12 +24,15 @@ 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; @@ -58,7 +61,10 @@ public void close() { } 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("dn_thrift_max_frame_size")); + 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(prefixLength, transport.writeBuffer.getPos()); assertEquals(0, underlying.length()); @@ -135,6 +141,67 @@ public void testCompressedFrameLimitAppliesAfterCompression() throws Exception { 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, Integer.MAX_VALUE, false); + Field position = AutoScalingBufferWriteTransport.class.getDeclaredField("pos"); + position.setAccessible(true); + int oldPosition = AutoResizingBuffer.MAX_ARRAY_SIZE - 4; + position.setInt(transport.writeBuffer, oldPosition); + 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(oldPosition, transport.writeBuffer.getPos()); + assertTrue(closed[0]); + } + } + @Test public void testInvalidCompressionBufferSizesCloseTransport() throws Exception { for (int size : new int[] {-1, Integer.MAX_VALUE}) { 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 index 071dc4fa6f57a..469bcb9e2c1b6 100644 --- 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 @@ -279,7 +279,9 @@ public void close() { false)); assertTrue(failure.getMessage(), failure.getMessage().contains("at least")); assertTrue(failure.getMessage(), failure.getMessage().contains("1024 bytes")); - assertTrue(failure.getMessage(), failure.getMessage().contains("dn_thrift_max_frame_size")); + 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()); From c7003fff631e252ae4f0796ae258124caffcf2a5 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Sat, 14 Nov 2026 09:47:10 +0800 Subject: [PATCH 7/7] Align RPC method ordering with style checks --- .../rpc/AutoScalingBufferWriteTransport.java | 24 +++++++++---------- .../rpc/TElasticFramedTransportTest.java | 2 +- 2 files changed, 13 insertions(+), 13 deletions(-) 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 45dfc1083968d..c193bd3ea3568 100644 --- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/AutoScalingBufferWriteTransport.java @@ -47,18 +47,6 @@ public void write(byte[] toWrite, int off, int len) throws TTransportException { write(toWrite, off, len, AutoResizingBuffer.MAX_ARRAY_SIZE); } - // 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 = buf.array(); - if (len > array.length - pos) { - return false; - } - System.arraycopy(toWrite, off, array, pos, len); - pos += len; - return true; - } - 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; @@ -78,6 +66,18 @@ void write(byte[] toWrite, int off, int len, int maxCapacity) throws TTransportE 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 = buf.array(); + if (len > array.length - pos) { + return false; + } + System.arraycopy(toWrite, off, array, pos, len); + pos += len; + return true; + } + public int getPos() { return pos; } 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 87b81ecbceb4b..d1873ede91ad0 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 @@ -190,7 +190,7 @@ public void close() { position.setAccessible(true); int oldPosition = AutoResizingBuffer.MAX_ARRAY_SIZE - 4; position.setInt(transport.writeBuffer, oldPosition); - byte[] buffer = transport.writeBuffer.getBuffer(); + final byte[] buffer = transport.writeBuffer.getBuffer(); TTransportException failure = assertThrows(TTransportException.class, () -> transport.write(new byte[16])); assertEquals(TTransportException.CORRUPTED_DATA, failure.getType());