From 72470f65a58ff75ea9ca90caf71a0181a85166cc Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 14 Oct 2026 12:11:34 +0800 Subject: [PATCH 01/15] Document exchange fan-in implementation scope --- ...2026-10-14-exchange-fanin-budget-design.md | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md new file mode 100644 index 000000000000..b2bc07d1ddbb --- /dev/null +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -0,0 +1,71 @@ + + +# Exchange 无序扇入与输入输出预算 + +基线：Apache master `aec1ad6139cb4dd1450725f53418cf40c26efb4e`。 +依据：[技术设计文档](https://timechor.feishu.cn/docx/EpDFdmS94ovMCWxluEOc3WGDn8d)，修订 861。 + +## 已确认范围 + +实现 FI 输入输出共享账目、Collect/MappingCollect 任意就绪消费，以及由无序汇集引起的计划属性修正。 +聚合使用 HashAggregation：只清空实际输入已失去连续性的 preGroupedSymbols；有效分支内的流式聚合保留。 +MergeSort、Join、ASOF、CrossJoin 只作依赖验证，不修改其消费算法，不新增进度保障机制。 +TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍需匹配实际输入。 + +## 额度与账目 + +F 为配置的单 FI 交换预算。INPUT=F/2，OUTPUT=F-INPUT；同 FI 的所有 pipeline 继承相同固定值。 +远端 SourceHandle 计入接收 FI 的 INPUT；SinkChannel 和 SharedTsBlockQueue 计入生产 FI 的 OUTPUT。 +本地共享队列只计费一次，创建先后不改变所有者。 + +MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。 +FI 两侧使用 AtomicLong；planNode 字节明细通过 compute 增减。 +申请、回滚、释放均对应修改全局池、FI 单侧与 planNode 三层账目。 +额度判断使用本次增加后的单侧总数；稳定快照中的总数等于同侧明细之和。 +保留现有最高优先级查询的零计费兜底与交换组件的收发、等待、释放机制。 + +## Collect 生命周期 + +AnyChildBlocked 缓存直接 child.isBlocked 的 Future，每个 pending 周期只有一个 listener。 +新轮次重定向现有 listener；发布后检查完成状态，覆盖注册与完成交错。 +回调只记录就绪或错误，在锁外完成选择 Future；所有 child 方法由 Driver 线程调用。 +重复 isBlocked 与输出背压保留选择；取消对外等待不向 child 传播。 + +一次 next 消耗一次选择。null/空块刷新等待，不表示结束。 +hasNext=false 后检查 isFinished，确认结束才关闭 child 并移除记录。 +MappingCollect 按实际 childId 映射。关闭唤醒待完成选择，忽略迟到回调，逐一尝试关闭全部输入。 +保留内存 R 为各存活 child 保留内存之和；峰值为 max(R,max(R-r_i+p_i))；返回大小为 max(b_i)。 + +## 计划属性 + +多路 Collect/Union 打断分支连续性和排序；单路可透传。 +Filter 与符号透传 Project 不恢复属性；有效 Sort/MergeSort 可以重新建立属性。 +在 SINGLE/PARTIAL/FINAL 实际 child 组装后校验聚合，包括 DISTINCT/mask、普通分阶段聚合及 Union 下推。 +Group 在分区连续性失效时改为完整 Sort。复用 master 的 StreamSort 前缀校验和 Window/Fill 修复。 +SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除恢复排序。 + +## 验证 + +覆盖共享额度并发、INPUT/OUTPUT 隔离、所有者、奇数 F、超大申请、回滚释放取消及稳定账目一致性； +覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射； +覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。 +运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。 From 19d4119f7deaaa638ad7f94c620de048f072772d Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 14 Oct 2026 12:47:59 +0800 Subject: [PATCH 02/15] Share fragment exchange budgets and consume ready Collect inputs --- ...2026-10-14-exchange-fanin-budget-design.md | 28 ++ .../it/query/recent/IoTDBCollectFanInIT.java | 276 +++++++++++ .../operator/process/AnyChildBlocked.java | 242 ++++++++++ .../operator/process/CollectOperator.java | 136 ++++-- .../process/MappingCollectOperator.java | 56 ++- .../operator/process/CollectOperatorTest.java | 442 ++++++++++++++++++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../exchange/SharedTsBlockQueue.java | 9 +- .../exchange/sink/ShuffleSinkHandle.java | 4 +- .../execution/exchange/sink/SinkChannel.java | 9 +- .../exchange/source/SourceHandle.java | 9 +- .../execution/memory/MemoryPool.java | 195 ++++++-- .../execution/memory/MemorySide.java | 31 ++ .../DataNodeTableOperatorGenerator.java | 4 +- .../planner/LocalExecutionPlanContext.java | 46 +- .../plan/planner/LocalExecutionPlanner.java | 4 +- .../plan/planner/OperatorTreeGenerator.java | 6 +- .../TableDistributedPlanGenerator.java | 122 ++++- .../optimizations/ParallelizeGrouping.java | 16 + .../optimizations/SortElimination.java | 25 + .../exchange/MPPDataExchangeManagerTest.java | 53 +++ .../execution/exchange/SourceHandleTest.java | 3 +- .../execution/memory/MemoryPoolTest.java | 115 ++++- .../LocalExecutionPlanContextBudgetTest.java | 78 ++++ .../planner/CollectInputPropertyTest.java | 149 ++++++ .../ParallelizeGroupingTest.java | 54 +++ 27 files changed, 1968 insertions(+), 148 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java create mode 100644 iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java create mode 100644 iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index b2bc07d1ddbb..a817039cab4b 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -69,3 +69,31 @@ SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除 覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射； 覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。 运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。 + +## 实现验证记录（2026-10-14） + +第一版实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。 +集成验证包含普通与 DISTINCT 聚合、Union 分阶段聚合、嵌套窗口、Fill、排序/TopK、Join 依赖、树模型以及树表视图。 + +构建先执行以下干净构建，生成当前分支的节点包并安装本地依赖： + +```sh +mvn -o clean install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests +``` + +全 reactor 中文编译与英文编译/回归均通过；英文 test 阶段包含完整 test-compile： + +```sh +mvn -o test-compile -DskipTests -P with-zh-locale +mvn -o test -Dtest=CollectOperatorTest,MemoryPoolTest,LocalExecutionPlanContextBudgetTest,CollectInputPropertyTest,SharedTsBlockQueueTest,SourceHandleTest,SinkChannelTest,LocalSourceHandleTest,LocalSinkChannelTest,MPPDataExchangeManagerTest,ParallelizeGroupingTest,TableFunctionSortTest -Dsurefire.failIfNoSpecifiedTests=false +``` + +两个集成环境均只运行本分支新增的测试类： + +```sh +mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableSimpleIT -P with-integration-tests +mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableClusterIT -P with-integration-tests +``` + +使用离线模式是为了避免远端 SNAPSHOT ZIP 下载停滞；节点执行的是当前分支重新打包的代码。 +Spotless、Checkstyle、新增文件许可证及中英文消息键名/占位符一致性检查通过。 diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java new file mode 100644 index 000000000000..adc72b6046e5 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -0,0 +1,276 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBCollectFanInIT { + private static final String DATABASE = "collect_fan_in"; + private static final int DEVICES = 24; + private static final int ROWS = 8; + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setMaxTsBlockLineNumber(2) + .setSortBufferSize(128 * 1024) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDataRegionPerDataNode(3); + EnvFactory.getEnv().initClusterEnvironment(); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE " + DATABASE); + statement.execute("USE " + DATABASE); + for (String table : new String[] {"readings1", "readings2"}) { + statement.execute( + "CREATE TABLE " + table + " (device STRING TAG, zone STRING TAG, value INT32 FIELD)"); + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + statement.execute( + "INSERT INTO " + + table + + " (time,device,zone,value) VALUES (" + + value + + ",'" + + device(device) + + "'," + + (zone(device) == null ? "NULL" : "'" + zone(device) + "'") + + "," + + value + + ")"); + } + } + } + statement.execute("FLUSH"); + } + } + + @Test + public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE root.collect_tree"); + for (int device = 0; device < DEVICES; device++) { + String path = "root.collect_tree.z" + device % 3 + "." + device(device); + statement.execute("CREATE ALIGNED TIMESERIES " + path + "(value INT32)"); + for (int value = 1; value <= ROWS; value++) { + statement.execute( + "INSERT INTO " + path + "(time,value) ALIGNED VALUES (" + value + "," + value + ")"); + } + } + try (ResultSet result = + statement.executeQuery("SELECT value FROM root.collect_tree.*.* ALIGN BY DEVICE")) { + int rows = 0; + while (result.next()) { + assertEquals(result.getLong("Time"), result.getLong("value")); + rows++; + } + assertEquals(DEVICES * ROWS, rows); + } + } + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + statement.execute( + "CREATE VIEW tree_readings(zone STRING TAG, device STRING TAG, value INT32 FIELD) AS root.collect_tree.**"); + try (ResultSet result = + statement.executeQuery( + "SELECT device, SUM(value) FROM (" + + "SELECT device, value FROM tree_readings UNION ALL SELECT device, value FROM tree_readings" + + ") u GROUP BY device ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(72.0, result.getDouble(2), 0.0); + } + assertFalse(result.next()); + } + } + } + + @AfterClass + public static void tearDown() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + private static String device(int id) { + return "d" + (id < 10 ? "0" : "") + id; + } + + private static String zone(int id) { + return id % 3 == 0 ? null : "z" + id % 3; + } + + @Test + public void testGroupedAggregationAfterGlobalWindow() throws Exception { + checkGroups( + "SELECT device, zone, value FROM (SELECT device, zone, value, " + + "COUNT(*) OVER () AS n FROM readings1) t WHERE n > 0", + 1); + } + + @Test + public void testUnionPartialFinalAndDistinctMask() throws Exception { + String input = + "SELECT device, zone, value FROM readings1 UNION ALL " + + "SELECT device, zone, value FROM readings2"; + checkGroups(input, 2); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + // Without DISTINCT, this query exercises branch PARTIAL and post-Collect FINAL aggregation. + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, SUM(value) FROM (" + + input + + ") u GROUP BY device, zone ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(72.0, result.getDouble(3), 0.0); + } + assertFalse(result.next()); + } + } + } + + @Test + public void testOrderedFanInAndJoinDependencies() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + for (String limit : new String[] {"", " LIMIT 13"}) { + try (ResultSet result = + statement.executeQuery( + "SELECT device, value FROM (SELECT device, value FROM readings1 UNION ALL " + + "SELECT device, value FROM readings2) u ORDER BY device, value" + + limit)) { + int expectedRows = limit.isEmpty() ? DEVICES * ROWS * 2 : 13; + for (int row = 0; row < expectedRows; row++) { + assertTrue(result.next()); + assertEquals(device(row / (ROWS * 2)), result.getString(1)); + assertEquals(row / 2 % ROWS + 1, result.getInt(2)); + } + assertFalse(result.next()); + } + } + try (ResultSet result = + statement.executeQuery( + "SELECT COUNT(*) FROM readings1 a JOIN readings2 b " + + "ON a.device = b.device AND a.time = b.time WHERE a.value = b.value")) { + assertTrue(result.next()); + assertEquals(DEVICES * ROWS, result.getLong(1)); + assertFalse(result.next()); + } + } + } + + private void checkGroups(String input, int copies) throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, SUM(value), " + + "COUNT(*), COUNT(DISTINCT value) FROM (" + + input + + ") u GROUP BY device, zone ORDER BY device")) { + for (int device = 0; device < DEVICES; device++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(36.0 * copies, result.getDouble(3), 0.0); + assertEquals((long) ROWS * copies, result.getLong(4)); + assertEquals(ROWS, result.getLong(5)); + } + assertFalse(result.next()); + } + } + } + + @Test + public void testNestedWindowAfterCollectRestoresCompletePartitions() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, " + + "SUM(value) OVER (PARTITION BY device, zone ORDER BY value) AS running_sum FROM (" + + "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1) t " + + "WHERE n > 0 ORDER BY device, value")) { + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(zone(device), result.getString(2)); + assertEquals(value, result.getInt(3)); + assertEquals((double) value * (value + 1) / 2, result.getDouble(4), 0.0); + } + } + assertFalse(result.next()); + } + } + } + + @Test + public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, COUNT(*) OVER () AS n " + + "FROM readings1 FILL METHOD CONSTANT 0 ORDER BY device, value DESC")) { + for (int device = 0; device < DEVICES; device++) { + for (int value = ROWS; value > 0; value--) { + assertTrue(result.next()); + assertEquals(device(device), result.getString(1)); + assertEquals(value, result.getInt(3)); + assertEquals((long) DEVICES * ROWS, result.getLong(4)); + } + } + assertFalse(result.next()); + } + } + } +} diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java new file mode 100644 index 000000000000..c0397bf1ff5a --- /dev/null +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java @@ -0,0 +1,242 @@ +/* + * 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.calc.execution.operator.process; + +import org.apache.iotdb.calc.execution.operator.Operator; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.utils.Accountable; +import org.apache.tsfile.utils.RamUsageEstimator; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; + +import static com.google.common.util.concurrent.Futures.getDone; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + +/** + * Observes direct children without calling operators from completion callbacks. The driver owns + * selection/invalidation; the lock protects listener records shared with asynchronous completions. + */ +final class AnyChildBlocked implements Accountable { + private static final long INSTANCE_SIZE = + RamUsageEstimator.shallowSizeOfInstance(AnyChildBlocked.class); + private static final long LISTENER_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ChildFutureListener.class); + private static final long ROUND_SIZE = + RamUsageEstimator.shallowSizeOfInstance(WaitRound.class) + + RamUsageEstimator.shallowSizeOfInstance(SettableFuture.class); + + private final List children; + private final Object lock = new Object(); + private final ChildFutureListener[] listeners; + private WaitRound currentRound; + private Throwable firstFailure; + private boolean closed; + private int nextChild; + + AnyChildBlocked(List children) { + this.children = children; + this.listeners = new ChildFutureListener[children.size()]; + } + + ListenableFuture getAnyChild() { + WaitRound round; + synchronized (lock) { + if (currentRound != null) { + return currentRound.selected; + } + round = new WaitRound(); + currentRound = round; + for (ChildFutureListener listener : listeners) { + if (listener != null) { + listener.target = round; + } + } + } + + // Activate every live direct child, even if an earlier one has already become ready. Only + // invalidated records call child.isBlocked(); a long-pending future keeps its one listener. + for (int offset = 0; offset < children.size(); offset++) { + int childId = (nextChild + offset) % children.size(); + Operator child; + ChildFutureListener listener; + synchronized (lock) { + if (closed) { + break; + } + child = children.get(childId); + listener = listeners[childId]; + } + if (child == null) { + continue; + } + if (listener == null) { + ListenableFuture observedFuture = child.isBlocked(); + synchronized (lock) { + if (closed) { + break; + } + listener = new ChildFutureListener(childId, observedFuture, round); + listeners[childId] = listener; + } + // The target is published before registration. addListener also handles a future that + // finishes between isBlocked and registration, including inline completion. + observedFuture.addListener(listener, directExecutor()); + } else if (listener.observedFuture.isDone()) { + // A completed record may have notified a previous round. Recheck it for this round. + listener.run(); + } + } + + Throwable failure; + boolean isClosed; + synchronized (lock) { + failure = firstFailure; + isClosed = closed; + } + if (isClosed || children.isEmpty()) { + round.selected.set(-1); + } else if (failure != null) { + round.selected.setException(failure); + } + return round.selected; + } + + void checkFailure() throws Exception { + Throwable failure; + synchronized (lock) { + failure = firstFailure; + } + if (failure instanceof Exception) { + throw (Exception) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw new RuntimeException(failure); + } + } + + void invalidateChild(int childId) { + synchronized (lock) { + listeners[childId] = null; + if (currentRound != null) { + currentRound.valid = false; + currentRound = null; + } + nextChild = (childId + 1) % children.size(); + } + } + + void finishChild(int childId) { + invalidateChild(childId); + } + + void close() { + WaitRound round; + synchronized (lock) { + closed = true; + round = currentRound; + currentRound = null; + if (round != null) { + round.valid = false; + } + for (ChildFutureListener listener : listeners) { + if (listener != null) { + listener.target = null; + } + } + Arrays.fill(listeners, null); + } + if (round != null) { + round.selected.set(-1); + } + } + + @Override + public long ramBytesUsed() { + synchronized (lock) { + long size = + INSTANCE_SIZE + + RamUsageEstimator.shallowSizeOfInstance(Object.class) + + RamUsageEstimator.shallowSizeOf(listeners); + for (ChildFutureListener listener : listeners) { + if (listener != null) { + size += LISTENER_SIZE; + } + } + return size + (currentRound == null ? 0 : ROUND_SIZE); + } + } + + private static final class WaitRound { + private final SettableFuture selected = SettableFuture.create(); + private boolean valid = true; + } + + private final class ChildFutureListener implements Runnable { + private final int childId; + private final ListenableFuture observedFuture; + private WaitRound target; + + private ChildFutureListener(int childId, ListenableFuture observedFuture, WaitRound target) { + this.childId = childId; + this.observedFuture = observedFuture; + this.target = target; + } + + @Override + public void run() { + Throwable failure = null; + try { + getDone(observedFuture); + } catch (ExecutionException e) { + failure = e.getCause(); + } catch (CancellationException e) { + failure = e; + } + WaitRound round; + synchronized (lock) { + if (closed || listeners[childId] != this) { + return; + } + if (failure != null && firstFailure == null) { + firstFailure = failure; + } + round = target; + if (round == null || !round.valid || round != currentRound) { + return; + } + failure = firstFailure; + } + // Completing a selection can invoke downstream listeners synchronously. Do not hold lock. + if (failure == null) { + round.selected.set(childId); + } else { + round.selected.setException(failure); + } + } + } +} diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java index 921d273d04e8..2cc5deec59ea 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java @@ -30,60 +30,102 @@ import java.util.List; +import static com.google.common.util.concurrent.Futures.getDone; +import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; + public class CollectOperator implements ProcessOperator { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class); private final CommonOperatorContext operatorContext; protected final List children; - private boolean inited = false; - - protected int currentIndex; + private final AnyChildBlocked anyChildBlocked; + private ListenableFuture selectionFuture; + private ListenableFuture exposedFuture; + private int remainingChildren; + private boolean closed; public CollectOperator(CommonOperatorContext operatorContext, List children) { this.operatorContext = operatorContext; this.children = children; - this.currentIndex = 0; + this.anyChildBlocked = new AnyChildBlocked(children); + this.remainingChildren = children.size(); } @Override - public boolean hasNext() throws Exception { - return currentIndex < children.size(); + public boolean hasNext() { + return !closed && remainingChildren > 0; } @Override public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - return children.get(currentIndex).nextWithTimer(); - } else { - closeCurrentChild(currentIndex); - currentIndex++; + anyChildBlocked.checkFailure(); + if (!hasNext()) { + return null; + } + if (selectionFuture == null) { + isBlocked(); + } + anyChildBlocked.checkFailure(); + if (!selectionFuture.isDone()) { + return null; + } + int childId = getDone(selectionFuture); + // Consume this selection before invoking any operator method. Another next must acquire a + // new readiness observation even when this advancement returns null or an empty block. + selectionFuture = null; + exposedFuture = null; + if (childId < 0 || closed || children.get(childId) == null) { return null; } + Operator child = children.get(childId); + try { + if (child.hasNextWithTimer()) { + TsBlock block = child.nextWithTimer(); + return block == null ? null : transformBlock(childId, block); + } + if (child.isFinished()) { + closeCurrentChild(childId); + } + return null; + } finally { + anyChildBlocked.invalidateChild(childId); + } + } + + protected TsBlock transformBlock(int childId, TsBlock block) { + return block; } protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); + Operator child = children.get(index); + if (child == null) { + return; + } + // Retire the record even if closing its subtree throws. + anyChildBlocked.finishChild(index); children.set(index, null); + remainingChildren--; + child.close(); } @Override public ListenableFuture isBlocked() { - if (!inited) { - inited = true; - for (Operator child : children) { - child.isBlocked(); - } - } - if (currentIndex >= children.size()) { + if (!hasNext()) { return NOT_BLOCKED; } - return children.get(currentIndex).isBlocked(); + if (selectionFuture == null) { + selectionFuture = anyChildBlocked.getAnyChild(); + } + if (exposedFuture == null || exposedFuture.isCancelled()) { + exposedFuture = nonCancellationPropagating(selectionFuture); + } + return exposedFuture; } @Override - public boolean isFinished() throws Exception { - return currentIndex >= children.size(); + public boolean isFinished() { + return !hasNext(); } @Override @@ -93,35 +135,67 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { - for (int i = currentIndex, n = children.size(); i < n; i++) { - Operator currentChild = children.get(i); - if (currentChild != null) { + if (closed) { + return; + } + closed = true; + anyChildBlocked.close(); + selectionFuture = null; + exposedFuture = null; + Exception firstException = null; + for (int i = 0; i < children.size(); i++) { + try { closeCurrentChild(i); + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } else if (firstException != e) { + firstException.addSuppressed(e); + } } } + if (firstException != null) { + throw firstException; + } } @Override public long calculateMaxPeekMemory() { - long maxPeekMemory = 0; + long retained = calculateRetainedSizeAfterCallingNext(); + long peak = retained; for (Operator child : children) { - maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + if (child != null) { + peak = + Math.max( + peak, + retained + - child.calculateRetainedSizeAfterCallingNext() + + child.calculateMaxPeekMemoryWithCounter()); + } } - return maxPeekMemory; + return peak; } @Override public long calculateMaxReturnSize() { long maxReturnSize = 0; for (Operator child : children) { - maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize()); + if (child != null) { + maxReturnSize = Math.max(maxReturnSize, child.calculateMaxReturnSize()); + } } return maxReturnSize; } @Override public long calculateRetainedSizeAfterCallingNext() { - return 0L; + long retained = 0; + for (Operator child : children) { + if (child != null) { + retained += child.calculateRetainedSizeAfterCallingNext(); + } + } + return retained; } @TestOnly @@ -132,7 +206,9 @@ public List getChildren() { @Override public long ramBytesUsed() { return INSTANCE_SIZE + + anyChildBlocked.ramBytesUsed() + children.stream() + .filter(java.util.Objects::nonNull) .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject) .sum() + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext); diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java index 8bc1c18e9e66..02850bf0caf1 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/MappingCollectOperator.java @@ -45,37 +45,43 @@ public MappingCollectOperator( List> mappings) { super(operatorContext, children); this.mappings = mappings; - outputColumnsCount = mappings.get(0).size(); + outputColumnsCount = mappings.isEmpty() ? 0 : mappings.get(0).size(); } @Override - public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - TsBlock tsBlock = children.get(currentIndex).nextWithTimer(); - if (tsBlock == null) { - return null; - } else { - Column[] columns = new Column[outputColumnsCount]; - List mapping = mappings.get(currentIndex); - for (int i = 0; i < columns.length; i++) { - columns[i] = tsBlock.getColumn(mapping.get(i)); - } - return TsBlock.wrapBlocksWithoutCopy( - tsBlock.getPositionCount(), - new RunLengthEncodedColumn( - CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()), - columns); - } - } else { - closeCurrentChild(currentIndex); - currentIndex++; - return null; + protected TsBlock transformBlock(int childId, TsBlock tsBlock) { + Column[] columns = new Column[outputColumnsCount]; + List mapping = mappings.get(childId); + for (int i = 0; i < columns.length; i++) { + columns[i] = tsBlock.getColumn(mapping.get(i)); } + return TsBlock.wrapBlocksWithoutCopy( + tsBlock.getPositionCount(), + new RunLengthEncodedColumn( + CommonOperatorUtils.TIME_COLUMN_TEMPLATE, tsBlock.getPositionCount()), + columns); } + @Override protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); - children.set(index, null); - mappings.set(index, null); + try { + super.closeCurrentChild(index); + } finally { + mappings.set(index, null); + } + } + + @Override + public long ramBytesUsed() { + long size = + super.ramBytesUsed() + + INSTANCE_SIZE + - RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class); + for (List mapping : mappings) { + if (mapping != null) { + size += RamUsageEstimator.shallowSizeOf(mapping.toArray()); + } + } + return size; } } diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java new file mode 100644 index 000000000000..abf7a82acf3d --- /dev/null +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -0,0 +1,442 @@ +/* + * 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.calc.execution.operator.process; + +import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; + +import com.google.common.util.concurrent.ForwardingListenableFuture; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.junit.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.Executor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +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.assertTrue; +import static org.junit.Assert.fail; + +public class CollectOperatorTest { + private static TsBlock block(int... values) { + TsBlockBuilder builder = + new TsBlockBuilder(Collections.nCopies(values.length, TSDataType.INT32)); + builder.getTimeColumnBuilder().writeLong(0); + for (int i = 0; i < values.length; i++) { + builder.getColumnBuilder(i).writeInt(values[i]); + } + builder.declarePosition(); + return builder.build(); + } + + private static CollectOperator collect(Operator... children) { + return new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children))); + } + + @Test + public void testReadyChildAndStableSelectionDuringBackpressure() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = collect(a, b); + ListenableFuture selected = collect.isBlocked(); + assertTrue(selected.isDone()); + a.ready(); + assertSame(selected, collect.isBlocked()); + assertSame(b.output, collect.next()); + assertEquals(0, a.advances); + assertEquals(1, b.advances); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + collect.close(); + } + + @Test + public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exception { + for (int count : new int[] {100, 1000}) { + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) { + children.add(new Child(block(i))); + } + CollectOperator collect = new CollectOperator(new TestContext(), children); + Child ready = (Child) children.get(count - 1); + for (int round = 0; round < 5; round++) { + ListenableFuture wait = collect.isBlocked(); + assertFalse(wait.isDone()); + for (int repeat = 0; repeat < 10; repeat++) { + assertSame(wait, collect.isBlocked()); + } + ready.ready(); + wait.get(5, TimeUnit.SECONDS); + assertSame(ready.output, collect.next()); + ready.reset(); + } + for (int i = 0; i < count - 1; i++) { + Child child = (Child) children.get(i); + assertEquals(1, child.waitCalls); + assertEquals(1, child.wait.registrations.get()); + } + collect.close(); + } + } + + @Test + public void testPendingListenerRetargetedAndStaleCompletionIgnored() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + CollectOperator collect = collect(a, b); + ListenableFuture first = collect.isBlocked(); + b.ready(); + assertSame(b.output, collect.next()); + Runnable stale = b.wait.listener; + b.reset(); + ListenableFuture second = collect.isBlocked(); + assertFalse(second.isDone()); + stale.run(); + assertFalse(second.isDone()); + a.ready(); + second.get(5, TimeUnit.SECONDS); + assertEquals(1, a.wait.registrations.get()); + assertSame(a.output, collect.next()); + assertTrue(first.isDone()); + collect.close(); + } + + @Test + public void testNullEmptyAndUnfinishedChildRefreshReadiness() throws Exception { + Child child = new Child(null); + child.ready(); + CollectOperator collect = collect(child); + assertTrue(collect.isBlocked().isDone()); + assertNull(collect.next()); + child.reset(); + assertFalse(collect.isBlocked().isDone()); + child.output = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)).build(); + child.ready(); + assertTrue(collect.next().isEmpty()); + child.reset(); + child.hasNext = false; + child.ready(); + assertNull(collect.next()); + assertFalse(collect.isFinished()); + assertEquals(0, child.closes); + child.finished = true; + assertNull(collect.next()); + assertTrue(collect.isFinished()); + assertEquals(1, child.closes); + assertTrue(collect.isBlocked().isDone()); + } + + @Test + public void testCancellationDoesNotPropagateToChild() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = collect(child); + ListenableFuture exposed = collect.isBlocked(); + assertTrue(exposed.cancel(true)); + assertFalse(child.wait.isCancelled()); + ListenableFuture replacement = collect.isBlocked(); + assertFalse(replacement.isCancelled()); + child.ready(); + replacement.get(5, TimeUnit.SECONDS); + assertSame(child.output, collect.next()); + collect.close(); + } + + @Test + public void testFailureOfUnselectedChildIsNotLost() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.ready(); + CollectOperator collect = collect(a, b); + assertTrue(collect.isBlocked().isDone()); + IOException expected = new IOException("child failure"); + b.wait.delegate.setException(expected); + try { + collect.next(); + fail("must propagate the failed child even when another child was selected"); + } catch (IOException e) { + assertSame(expected, e); + } finally { + collect.close(); + } + assertEquals(0, a.advances); + } + + @Test + public void testChildCancellationPropagates() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = collect(child); + ListenableFuture wait = collect.isBlocked(); + child.wait.cancel(false); + assertTrue(wait.isDone()); + try { + collect.next(); + fail("must propagate child cancellation"); + } catch (CancellationException expected) { + assertEquals(0, child.advances); + } finally { + collect.close(); + } + } + + @Test + public void testCompletionDuringRegistrationAndOnAnotherThread() throws Exception { + Child a = new Child(block(1)); + a.wait.beforeRegistration = a::ready; + CollectOperator collect = collect(a); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + a.reset(); + ListenableFuture pending = collect.isBlocked(); + Thread completion = new Thread(a::ready); + completion.start(); + pending.get(5, TimeUnit.SECONDS); + completion.join(); + assertSame(a.output, collect.next()); + collect.close(); + } + + @Test + public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.closeFailure = new IOException("close failure"); + CollectOperator collect = collect(a, b); + ListenableFuture pending = collect.isBlocked(); + try { + collect.close(); + fail("must propagate close failure"); + } catch (IOException e) { + assertSame(a.closeFailure, e); + } + pending.get(5, TimeUnit.SECONDS); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + a.ready(); + b.ready(); + assertTrue(collect.isFinished()); + assertNull(collect.next()); + collect.close(); + assertEquals(1, b.closes); + } + + @Test + public void testMappingUsesActualChildIdentity() throws Exception { + Child a = new Child(block(10, 11)); + Child b = new Child(block(20, 21)); + b.ready(); + List> mappings = new ArrayList<>(); + mappings.add(Collections.singletonList(0)); + mappings.add(Collections.singletonList(1)); + MappingCollectOperator collect = + new MappingCollectOperator( + new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings); + assertTrue(collect.isBlocked().isDone()); + assertEquals(21, collect.next().getColumn(0).getInt(0)); + a.ready(); + assertEquals(10, collect.next().getColumn(0).getInt(0)); + collect.close(); + assertNull(mappings.get(0)); + assertNull(mappings.get(1)); + } + + @Test + public void testMemoryIncludesAllLiveRetainedStates() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.retained = 30; + a.peak = 100; + a.returnSize = 10; + b.retained = 50; + b.peak = 80; + b.returnSize = 20; + CollectOperator collect = collect(a, b); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); + assertEquals(20, collect.calculateMaxReturnSize()); + assertTrue(collect.ramBytesUsed() > 0); + a.hasNext = false; + a.finished = true; + a.ready(); + collect.next(); + assertEquals(50, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(80, collect.calculateMaxPeekMemory()); + collect.close(); + assertEquals(0, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateMaxReturnSize()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + } + + private static final class CountingFuture + extends ForwardingListenableFuture.SimpleForwardingListenableFuture { + private final SettableFuture delegate; + private final AtomicInteger registrations = new AtomicInteger(); + private Runnable beforeRegistration; + private Runnable listener; + + private CountingFuture() { + this(SettableFuture.create()); + } + + private CountingFuture(SettableFuture delegate) { + super(delegate); + this.delegate = delegate; + } + + @Override + public void addListener(Runnable listener, Executor executor) { + this.listener = listener; + registrations.incrementAndGet(); + if (beforeRegistration != null) { + beforeRegistration.run(); + } + super.addListener(listener, executor); + } + } + + private static final class TestContext extends CommonOperatorContext { + private TestContext() { + super(0, new PlanNodeId("collect"), "test"); + } + + @Override + public MemoryReservationManager getMemoryReservationContext() { + return null; + } + + @Override + public int getFragmentId() { + return 0; + } + + @Override + public int getPipelineId() { + return 0; + } + + @Override + public long ramBytesUsed() { + return 0; + } + } + + private static final class Child implements Operator { + private final Thread driver = Thread.currentThread(); + private final CommonOperatorContext context = new TestContext(); + private CountingFuture wait = new CountingFuture(); + private TsBlock output; + private boolean hasNext = true; + private boolean finished; + private int advances; + private int waitCalls; + private int closes; + private IOException closeFailure; + private long retained, peak, returnSize; + + private Child(TsBlock output) { + this.output = output; + } + + private void ready() { + wait.delegate.set(null); + } + + private void reset() { + wait = new CountingFuture(); + } + + private void checkDriver() { + assertSame(driver, Thread.currentThread()); + } + + @Override + public CommonOperatorContext getOperatorContext() { + return context; + } + + @Override + public ListenableFuture isBlocked() { + checkDriver(); + waitCalls++; + return wait; + } + + @Override + public boolean hasNext() { + checkDriver(); + return hasNext; + } + + @Override + public TsBlock next() { + checkDriver(); + advances++; + return output; + } + + @Override + public boolean isFinished() { + checkDriver(); + return finished; + } + + @Override + public void close() throws Exception { + checkDriver(); + closes++; + if (closeFailure != null) throw closeFailure; + } + + @Override + public long calculateMaxPeekMemory() { + return peak; + } + + @Override + public long calculateMaxReturnSize() { + return returnSize; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return retained; + } + + @Override + public long ramBytesUsed() { + return 0; + } + } +} 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 b030af64cfa7..be8c3ade6477 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 @@ -3941,4 +3941,6 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "Only InMemoryDeviceEntryDataSet supports get inline device entries"; + public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "Plan node %s is already registered on memory side %s, not %s"; + public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "Memory reservation is not registered for query %s, fragment %s, plan node %s"; } 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 5a504e2a0175..13fa4222a40d 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 @@ -4697,4 +4697,6 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_ONLY_INMEMORYDEVICEENTRYDATASET_SUPPORTS_GET_INLINE_DEVICE_ENTRIES_07A52CAB = "只有 InMemoryDeviceEntryDataSet 支持获取内存中的设备条目"; + public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "PlanNode %s 已注册到内存侧 %s，不能注册到 %s"; + public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "查询 %s、Fragment %s、PlanNode %s 尚未注册内存账目"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java index 5812d4e37cdd..cb8c681e4e9d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java @@ -27,6 +27,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.utils.CommonUtils; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -85,7 +86,8 @@ public class SharedTsBlockQueue { private LocalSinkChannel sinkChannel; private long maxBytesCanReserve = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + MemorySide.OUTPUT.getMaxBytes( + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); private final boolean isHighestPriority; private volatile Throwable abortedCause = null; @@ -265,7 +267,10 @@ public ListenableFuture add(TsBlock tsBlock) { localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.OUTPUT); alreadyRegistered = true; } MemoryReservationResult reserveResult = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java index e54062158a3f..1de3ccb897df 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java @@ -22,6 +22,7 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -76,7 +77,8 @@ public class ShuffleSinkHandle implements ISinkHandle { /** max bytes this ShuffleSinkHandle can reserve. */ private long maxBytesCanReserve = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + MemorySide.OUTPUT.getMaxBytes( + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(ShuffleSinkHandle.class) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index d24a1552a5b2..0dec1cbf07f2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -30,6 +30,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SinkListener; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet; import org.apache.iotdb.db.utils.SetThreadName; @@ -131,7 +132,8 @@ private TsBlockInfo(TsBlock tsBlock, ByteBuffer serializedTsBlock, long tsBlockS /** max bytes this SinkChannel can reserve. */ private long maxBytesCanReserve = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + MemorySide.OUTPUT.getMaxBytes( + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); private final boolean isHighestPriority; @@ -235,7 +237,10 @@ public SinkChannel( localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.OUTPUT); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index aed3f9533fb1..b1dc2781ed65 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.db.queryengine.metric.DataExchangeCountMetricSet; import org.apache.iotdb.db.utils.SetThreadName; @@ -117,7 +118,8 @@ public class SourceHandle implements ISourceHandle { /** max bytes this SourceHandle can reserve. */ private long maxBytesCanReserve = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + MemorySide.INPUT.getMaxBytes( + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); /** * this is set to true after calling isBlocked() at least once which indicates that this @@ -223,7 +225,10 @@ public SourceHandle( localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, + fullFragmentInstanceId, + localPlanNodeId, + MemorySide.INPUT); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index b2d35568fa31..56c9974764a5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -40,6 +40,7 @@ import java.util.Queue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; import static com.google.common.util.concurrent.Futures.immediateVoidFuture; @@ -57,8 +58,8 @@ public static class MemoryReservationFuture extends AbstractFuture { private final long bytesToReserve; /** - * MemoryReservationFuture is created when SinkHandle or SourceHandle tries to reserve memory - * from pool. This field is max Bytes that SinkHandle or SourceHandle can reserve. + * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the + * pool. This limit is shared by all plan nodes registered on the same side of that fragment. */ private final long maxBytesCanReserve; @@ -154,8 +155,32 @@ public long getReservedBytes() { private final IMemoryBlock memoryBlock; private final long maxBytesPerFragmentInstance; - /** queryId -> fragmentInstanceId -> planNodeId -> bytesReserved. */ - private final Map>> queryMemoryReservations = + private static class PlanNodeReservation { + private final MemorySide side; + private volatile long bytesReserved; + + private PlanNodeReservation(MemorySide side) { + this.side = side; + } + + @Override + public String toString() { + return side + ":" + bytesReserved; + } + } + + private static class FragmentMemoryReservation { + private final Map planNodes = new ConcurrentHashMap<>(); + private final AtomicLong inputUsedBytes = new AtomicLong(); + private final AtomicLong outputUsedBytes = new AtomicLong(); + + private AtomicLong getUsedBytes(MemorySide side) { + return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; + } + } + + /** queryId -> fragmentInstanceId -> FI side totals and planNode details. */ + private final Map> queryMemoryReservations = new ConcurrentHashMap<>(); private final Queue> memoryReservationFutures = @@ -204,12 +229,27 @@ public int getMemoryReservationSize() { * planNodeId to queryMemoryReservationsMap first. */ public void registerPlanNodeIdToQueryMemoryMap( - String queryId, String fragmentInstanceId, String planNodeId) { + String queryId, String fragmentInstanceId, String planNodeId, MemorySide side) { + java.util.Objects.requireNonNull(side); synchronized (queryMemoryReservations) { queryMemoryReservations .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) - .computeIfAbsent(fragmentInstanceId, x -> new ConcurrentHashMap<>()) - .putIfAbsent(planNodeId, 0L); + .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()) + .planNodes + .compute( + planNodeId, + (key, reservation) -> { + if (reservation != null && reservation.side != side) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625, + planNodeId, + reservation.side, + side)); + } + return reservation == null ? new PlanNodeReservation(side) : reservation; + }); } } @@ -224,15 +264,19 @@ public void registerPlanNodeIdToQueryMemoryMap( */ public void deRegisterFragmentInstanceFromQueryMemoryMap( String queryId, String fragmentInstanceId, boolean forceDeregister) { - Map> queryRelatedMemory = queryMemoryReservations.get(queryId); + Map queryRelatedMemory = + queryMemoryReservations.get(queryId); if (queryRelatedMemory != null) { - Map fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); + FragmentMemoryReservation fragmentRelatedMemory = queryRelatedMemory.get(fragmentInstanceId); boolean hasPotentialMemoryLeak = false; // fragmentRelatedMemory could be null if the FI has not reserved any memory(For example, // next() of root operator returns no data) if (fragmentRelatedMemory != null) { hasPotentialMemoryLeak = - fragmentRelatedMemory.values().stream().anyMatch(value -> value != 0L); + fragmentRelatedMemory.inputUsedBytes.get() != 0L + || fragmentRelatedMemory.outputUsedBytes.get() != 0L + || fragmentRelatedMemory.planNodes.values().stream() + .anyMatch(value -> value.bytesReserved != 0L); } if (!forceDeregister && hasPotentialMemoryLeak) { // If hasPotentialMemoryLeak is true, it means that LocalSourceChannel/LocalSourceHandles @@ -252,9 +296,9 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( } if (hasPotentialMemoryLeak) { // hasPotentialMemoryLeak means that fragmentRelatedMemory is not null - List> invalidEntryList = - fragmentRelatedMemory.entrySet().stream() - .filter(entry -> entry.getValue() != 0L) + List> invalidEntryList = + fragmentRelatedMemory.planNodes.entrySet().stream() + .filter(entry -> entry.getValue().bytesReserved != 0L) .collect(Collectors.toList()); throw new MemoryLeakException( String.format( @@ -390,25 +434,26 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l Validate.notNull(queryId, DataNodeQueryMessages.EXCEPTION_QUERYID_CAN_NOT_BE_NULL_DOT_16639DBE); Validate.isTrue(bytes > 0L); - try { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .computeIfPresent( - planNodeId, - (k, reservedMemory) -> { - if (reservedMemory < bytes) { - throw new IllegalArgumentException( - DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); - } - return reservedMemory - bytes; - }); - } catch (NullPointerException e) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + if (fragment == null) { throw new IllegalArgumentException( DataNodeQueryMessages - .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2, - e); + .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2); } + fragment.planNodes.compute( + planNodeId, + (key, reservation) -> { + requireReservation(reservation, queryId, fragmentInstanceId, planNodeId); + if (reservation.bytesReserved < bytes) { + throw new IllegalArgumentException( + DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); + } + reservation.bytesReserved -= bytes; + fragment.getUsedBytes(reservation.side).addAndGet(-bytes); + return reservation; + }); memoryBlock.release(bytes); @@ -440,16 +485,34 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l } public long getQueryMemoryReservedBytes(String queryId) { - if (!queryMemoryReservations.containsKey(queryId)) { + Map fragments = queryMemoryReservations.get(queryId); + if (fragments == null) { return 0L; } long sum = 0; - for (Map map : queryMemoryReservations.get(queryId).values()) { - sum = sum + map.values().stream().reduce(0L, Long::sum); + for (FragmentMemoryReservation fragment : fragments.values()) { + sum += fragment.inputUsedBytes.get() + fragment.outputUsedBytes.get(); } return sum; } + public long getFragmentMemoryReservedBytes( + String queryId, String fragmentInstanceId, MemorySide side) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + return fragment == null ? 0L : fragment.getUsedBytes(side).get(); + } + + public long getPlanNodeMemoryReservedBytes( + String queryId, String fragmentInstanceId, String planNodeId) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + PlanNodeReservation reservation = fragment == null ? null : fragment.planNodes.get(planNodeId); + return reservation == null ? 0L : reservation.bytesReserved; + } + public long getReservedBytes() { return memoryBlock.getUsedMemoryInBytes(); } @@ -460,22 +523,68 @@ public boolean tryReserve( String planNodeId, long bytesToReserve, long maxBytesCanReserve) { + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve); - long queryRemainingBytes = - maxBytesCanReserve - - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, bytesToReserve, Long::sum); - return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() && queryRemainingBytes >= 0; + long sideUsedBytes = fragment.getUsedBytes(reservation.side).addAndGet(bytesToReserve); + fragment.planNodes.compute( + planNodeId, + (key, value) -> { + value.bytesReserved += bytesToReserve; + return value; + }); + return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() + && sideUsedBytes <= maxBytesCanReserve; } private void rollbackReserve( String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) { - queryMemoryReservations - .get(queryId) - .get(fragmentInstanceId) - .merge(planNodeId, -bytesToReserve, Long::sum); + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + fragment.planNodes.compute( + planNodeId, + (key, reservation) -> { + reservation.bytesReserved -= bytesToReserve; + fragment.getUsedBytes(reservation.side).addAndGet(-bytesToReserve); + return reservation; + }); memoryBlock.release(bytesToReserve); } + + private FragmentMemoryReservation getFragmentReservation( + String queryId, String fragmentInstanceId, String planNodeId) { + Map fragments = queryMemoryReservations.get(queryId); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(fragmentInstanceId); + if (fragment == null) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F, + queryId, + fragmentInstanceId, + planNodeId)); + } + return fragment; + } + + private PlanNodeReservation requireReservation( + PlanNodeReservation reservation, + String queryId, + String fragmentInstanceId, + String planNodeId) { + if (reservation == null) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages + .EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F, + queryId, + fragmentInstanceId, + planNodeId)); + } + return reservation; + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java new file mode 100644 index 000000000000..a59ef20d2684 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java @@ -0,0 +1,31 @@ +/* + * 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.execution.memory; + +/** The exchange budget charged by a registered plan node. */ +public enum MemorySide { + INPUT, + OUTPUT; + + public long getMaxBytes(long fragmentExchangeMaxBytes) { + long inputMaxBytes = fragmentExchangeMaxBytes / 2; + return this == INPUT ? inputMaxBytes : fragmentExchangeMaxBytes - inputMaxBytes; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java index cb2f6f02acb5..fd539be7467e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java @@ -343,7 +343,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte Operator child = node.getChildren().get(0).accept(this, context); List children = new ArrayList<>(1); children.add(child); - sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); context.getDriverContext().setSink(sinkHandle); return new IdentitySinkOperator( operatorContext, children, downStreamChannelIndex, sinkHandle); @@ -385,7 +385,7 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext if (!isSameNode) { context.addExchangeSumNum(1); } - sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index 7a5f52749aa4..bfed28ccc4ca 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -32,8 +32,10 @@ import org.apache.iotdb.db.queryengine.execution.driver.DriverContext; import org.apache.iotdb.db.queryengine.execution.driver.SchemaDriverContext; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; +import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TemplatedInfo; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; @@ -47,8 +49,6 @@ import org.apache.tsfile.read.filter.basic.Filter; import org.apache.tsfile.utils.Binary; import org.apache.tsfile.utils.Pair; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import javax.annotation.Nullable; @@ -70,7 +70,6 @@ // Attention: We should use thread-safe data structure for members that are shared by all pipelines public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext { - private static final Logger LOGGER = LoggerFactory.getLogger(LocalExecutionPlanContext.class); // Save operators in this pipeline, a new one will be created when creating another pipeline private final DriverContext driverContext; private final AtomicInteger nextOperatorId; @@ -83,6 +82,8 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext private List pipelineDriverFactories; private List exchangeOperatorList = new ArrayList<>(); private int exchangeSumNum = 0; + private final long inputMaxBytes; + private final long outputMaxBytes; private List cachedDataTypes; @@ -110,6 +111,10 @@ public LocalExecutionPlanContext( TypeProvider typeProvider, FragmentInstanceContext instanceContext, DataNodeQueryContext dataNodeQueryContext) { + long fragmentExchangeMaxBytes = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); + this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); this.typeProvider = typeProvider; this.allSensorsMap = new ConcurrentHashMap<>(); this.nextOperatorId = new AtomicInteger(0); @@ -121,6 +126,8 @@ public LocalExecutionPlanContext( // For creating subContext, differ from parent context mainly in driver context public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) { + this.inputMaxBytes = parentContext.inputMaxBytes; + this.outputMaxBytes = parentContext.outputMaxBytes; this.nextOperatorId = parentContext.nextOperatorId; this.typeProvider = parentContext.typeProvider; this.allSensorsMap = parentContext.allSensorsMap; @@ -142,6 +149,10 @@ public LocalExecutionPlanContext( final TypeProvider typeProvider, final FragmentInstanceContext instanceContext, final ISchemaRegion schemaRegion) { + long fragmentExchangeMaxBytes = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); + this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); this.allSensorsMap = new ConcurrentHashMap<>(); this.typeProvider = typeProvider; this.nextOperatorId = new AtomicInteger(0); @@ -247,10 +258,12 @@ public void setExchangeSumNum(int exchangeSumNum) { this.exchangeSumNum = exchangeSumNum; } - public long getMaxBytesOneHandleCanReserve() { - long maxBytesPerFI = - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); - return exchangeSumNum == 0 ? maxBytesPerFI : maxBytesPerFI / exchangeSumNum; + public long getInputMaxBytes() { + return inputMaxBytes; + } + + public long getOutputMaxBytes() { + return outputMaxBytes; } public void addExchangeSumNum(int addValue) { @@ -261,16 +274,16 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { this.exchangeOperatorList.add(exchangeOperator); } - public void setMaxBytesOneHandleCanReserve() { - long maxBytesOneHandleCanReserve = getMaxBytesOneHandleCanReserve(); - LOGGER.debug( - DataNodeQueryMessages - .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG, - maxBytesOneHandleCanReserve, - exchangeSumNum); + public void setExchangeMemoryLimits() { exchangeOperatorList.forEach( - exchangeOperator -> - exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve)); + exchangeOperator -> { + exchangeOperator + .getSourceHandle() + .setMaxBytesCanReserve( + exchangeOperator.getSourceHandle() instanceof PipelineSourceHandle + ? outputMaxBytes + : inputMaxBytes); + }); } public Set getAllSensors(IDeviceID deviceId, String sensorId) { @@ -303,6 +316,7 @@ public void setISink(ISink sink) { checkArgument( driverContext.getSink() == null, DataNodeQueryMessages.EXCEPTION_THERE_MUST_BE_AT_MOST_ONE_SINKNODE_A965AFE7); + sink.setMaxBytesCanReserve(outputMaxBytes); driverContext.setSink(sink); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index feb6b0cb74af..1a85cb148c4f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -133,7 +133,7 @@ public List plan( context.getTimePartitions().ifPresent(instanceContext::setTimePartitions); // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setMaxBytesOneHandleCanReserve(); + context.setExchangeMemoryLimits(); return context.getPipelineDriverFactories(); } @@ -166,7 +166,7 @@ public List plan( context.addPipelineDriverFactory(root, context.getDriverContext(), 0); // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setMaxBytesOneHandleCanReserve(); + context.setExchangeMemoryLimits(); return context.getPipelineDriverFactories(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java index b6a57db5445d..d643ec03752e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java @@ -2577,7 +2577,7 @@ public Operator visitExchange(ExchangeNode node, LocalExecutionPlanContext conte if (!isSameNode) { context.addExchangeSumNum(1); } - sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); @@ -2616,7 +2616,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte node.getPlanNodeId().getId(), context.getInstanceContext()); List children = dealWithConsumeChildrenOneByOneNode(node, context); - sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); context.getDriverContext().setSink(sinkHandle); return new IdentitySinkOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); @@ -2658,7 +2658,7 @@ public Operator visitShuffleSink(ShuffleSinkNode node, LocalExecutionPlanContext localInstanceId.toThrift(), node.getPlanNodeId().getId(), context.getInstanceContext()); - sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); context.getDriverContext().setSink(sinkHandle); return new ShuffleHelperOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java index 66f6da892f3e..d7461431824e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java @@ -630,11 +630,27 @@ public List visitTopK(TopKNode node, PlanContext context) { if (canTopKEliminated(node.getOrderingScheme(), node.getCount(), childrenNodes.get(0))) { return childrenNodes; } - node.setChildren(Collections.singletonList(childrenNodes.get(0))); - return Collections.singletonList(node); + TopKNode physical = + new TopKNode( + node.getPlanNodeId(), + childrenNodes, + node.getOrderingScheme(), + node.getCount(), + node.getOutputSymbols(), + canSortEliminated( + node.getOrderingScheme(), + nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()))); + physical.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); + return Collections.singletonList(physical); } - TopKNode newTopKNode = (TopKNode) node.clone(); + TopKNode newTopKNode = + new TopKNode( + node.getPlanNodeId(), + node.getOrderingScheme(), + node.getCount(), + node.getOutputSymbols(), + true); // root TopK is not a runtime filter producer. newTopKNode.setTopKRuntimeFilterSourceId(null); for (PlanNode child : childrenNodes) { @@ -649,7 +665,8 @@ public List visitTopK(TopKNode node, PlanContext context) { node.getOrderingScheme(), node.getCount(), node.getOutputSymbols(), - node.isChildrenDataInOrder()); + canSortEliminated( + node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId()))); regionTopK.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); newChild = regionTopK; } @@ -689,6 +706,12 @@ public List visitGroup(GroupNode node, PlanContext context) { for (PlanNode child : childrenNodes) { if (canSortEliminated(node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId()))) { result.add(child); + } else if (groupingInvalidatedByCollect( + child, node.getOrderingScheme().getOrderBy().subList(0, node.getPartitionKeyCount()))) { + SortNode sort = + new SortNode(queryId.genPlanNodeId(), child, node.getOrderingScheme(), false, false); + nodeOrderingMap.put(sort.getPlanNodeId(), sort.getOrderingScheme()); + result.add(sort); } else { GroupNode subGroupNode = new GroupNode( @@ -1100,7 +1123,7 @@ public List visitExternalTsFileAggregationScan( if (needFinalAggregation) { OrderingScheme childOrdering = nodeOrderingMap.get(result.get(0).getPlanNodeId()); finalAggregation.setChild(mergeChildrenViaCollectOrMergeSort(childOrdering, result)); - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -2045,7 +2068,7 @@ public List visitAggregation(AggregationNode node, PlanContext context } List childrenNodes = node.getChild().accept(this, context); - OrderingScheme childOrdering = nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()); + OrderingScheme childOrdering = commonOrdering(childrenNodes); if (node.isStreamable()) { // Child has Ordering, we need to check if it is the Ordering we expected if (childOrdering != null) { @@ -2118,7 +2141,7 @@ public List visitAggregation(AggregationNode node, PlanContext context if (physicalAggregation.isStreamable() && childOrdering != null) { nodeOrderingMap.put(planNodeId, expectedOrderingSchema); } - return aggregationNode; + return validateAggregationGrouping(aggregationNode); }) .collect(Collectors.toList()); @@ -2128,10 +2151,10 @@ public List visitAggregation(AggregationNode node, PlanContext context new CollectNode(queryId.genPlanNodeId(), aggregationNodes.get(0).getOutputSymbols()); collectNode.setChildren(aggregationNodes); splitResult.left.setChild(collectNode); - return Collections.singletonList(splitResult.left); + return Collections.singletonList(validateAggregationGrouping(splitResult.left)); } - return Collections.singletonList(physicalAggregation); + return Collections.singletonList(validateAggregationGrouping(physicalAggregation)); } // We cannot do multi-stage Aggregate if any aggregation-function is distinct. @@ -2144,7 +2167,9 @@ public List visitAggregation(AggregationNode node, PlanContext context nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes); node.setChild(physicalChild); return Collections.singletonList( - withRateFunctionInputOrdering(node, nodeOrderingMap.get(physicalChild.getPlanNodeId()))); + validateAggregationGrouping( + withRateFunctionInputOrdering( + node, nodeOrderingMap.get(physicalChild.getPlanNodeId())))); } Pair splitResult = split(node, symbolAllocator, queryId); AggregationNode intermediate = splitResult.right; @@ -2167,13 +2192,73 @@ public List visitAggregation(AggregationNode node, PlanContext context if (node.isStreamable() && childOrdering != null) { nodeOrderingMap.put(planNodeId, expectedOrderingSchema); } - return aggregationNode; + return validateAggregationGrouping(aggregationNode); }) .collect(Collectors.toList()); splitResult.left.setChild( mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()), childrenNodes)); - return Collections.singletonList(splitResult.left); + return Collections.singletonList(validateAggregationGrouping(splitResult.left)); + } + + private AggregationNode validateAggregationGrouping(AggregationNode node) { + if (!node.getPreGroupedSymbols().isEmpty() + && groupingInvalidatedByCollect(node.getChild(), node.getPreGroupedSymbols())) { + node.setPreGroupedSymbols(ImmutableList.of()); + nodeOrderingMap.remove(node.getPlanNodeId()); + } + return node; + } + + /** Whether an unordered fan-in invalidates a previously established grouping property. */ + private boolean groupingInvalidatedByCollect(PlanNode node, List grouping) { + if (grouping.isEmpty()) { + return false; + } + if ((node instanceof CollectNode || node instanceof UnionNode) + && node.getChildren().size() > 1) { + return true; + } + OrderingScheme ordering = nodeOrderingMap.get(node.getPlanNodeId()); + if (ordering != null && prefixMatched(ordering, grouping)) { + return false; + } + // A completed aggregation emits each full grouping key once, even when it hashes its input. + if (node instanceof AggregationNode) { + AggregationNode aggregation = (AggregationNode) node; + if (!aggregation.getStep().isOutputPartial() + && aggregation.getGroupingSetCount() == 1 + && grouping.containsAll(aggregation.getGroupingKeys())) { + return false; + } + } + if (node instanceof ProjectNode) { + ProjectNode project = (ProjectNode) node; + List inputs = new ArrayList<>(); + for (Symbol symbol : grouping) { + Expression expression = project.getAssignments().get(symbol); + if (!(expression instanceof SymbolReference)) { + // A computed key cannot inherit a branch-local continuity proof. + return containsUnorderedCollect(project.getChild()); + } + inputs.add(Symbol.from(expression)); + } + return groupingInvalidatedByCollect(project.getChild(), inputs); + } + for (PlanNode child : node.getChildren()) { + if (groupingInvalidatedByCollect(child, grouping)) { + return true; + } + } + return false; + } + + private boolean containsUnorderedCollect(PlanNode node) { + if ((node instanceof CollectNode || node instanceof UnionNode) + && node.getChildren().size() > 1) { + return true; + } + return node.getChildren().stream().anyMatch(this::containsUnorderedCollect); } private static AggregationNode withRateFunctionInputOrdering( @@ -2301,7 +2386,8 @@ public List visitAggregationTableScan( throw new IllegalStateException( DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS); } - resultTableScanNodeList = Collections.singletonList(finalAggregation); + resultTableScanNodeList = + Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return resultTableScanNodeList; @@ -2460,7 +2546,7 @@ private List constructSpilledAggregationTableScanByRegionReplicaSet( mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); } - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -2677,7 +2763,8 @@ public List visitAggregationTreeDeviceViewScan( throw new IllegalStateException( DataNodeQueryMessages.LIST_PLANNODE_SIZE_SHOULD_1_BUT_NOW_IS); } - resultTableScanNodeList = Collections.singletonList(finalAggregation); + resultTableScanNodeList = + Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return resultTableScanNodeList; @@ -2863,7 +2950,7 @@ private List constructSpilledAggregationTreeDeviceViewScanByRegionRepl mergeChildrenViaCollectOrMergeSort( nodeOrderingMap.get(result.get(0).getPlanNodeId()), result)); } - return Collections.singletonList(finalAggregation); + return Collections.singletonList(validateAggregationGrouping(finalAggregation)); } return result; } @@ -4007,6 +4094,9 @@ public NodeDistribution getNodeDistribution(PlanNodeId nodeId) { public void clearExpectedOrderingScheme() { expectedOrderingScheme = null; hasSortProperty = false; + // Grouping pushdown depends on this ordering. A boundary such as a global window + // must not pass an outer Group's partitioning request to its scans. + pushDownGrouping = false; } public void setExpectedOrderingScheme(OrderingScheme expectedOrderingScheme) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java index 66c4eff6e3e0..ac3fc65fe668 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGrouping.java @@ -34,6 +34,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis; import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode; @@ -101,6 +102,21 @@ public PlanNode visitPlan(PlanNode node, Context context) { return newNode; } + @Override + public PlanNode visitUnion(UnionNode node, Context context) { + if (node.getChildren().size() <= 1) { + return visitPlan(node, context); + } + if (!context.canSkip()) { + context.canOptimized = TO_SORT; + } + // Neither branch may establish a property for the interleaved Union output. + return node.replaceChildren( + node.getChildren().stream() + .map(child -> child.accept(this, new Context(null, 0))) + .collect(Collectors.toList())); + } + /** We need to make sure: context#partitionKey can match the prefix of childOrderSchema */ private void checkPrefixMatch(Context context, List childOrder) { if (context.canSkip()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java index 64ef1b8c2adc..c35f66128052 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/SortElimination.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.FillNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GapFillNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.PatternRecognitionNode; @@ -31,6 +32,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TableFunctionProcessorNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKRankingNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.relational.planner.iterative.rule.PruneTableScanColumns; @@ -38,6 +40,7 @@ import com.google.common.collect.ImmutableSet; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Optional; @@ -75,6 +78,28 @@ public PlanNode visitPlan(PlanNode node, Context context) { return newNode; } + @Override + public PlanNode visitCollect(CollectNode node, Context context) { + return visitUnorderedFanIn(node, context); + } + + @Override + public PlanNode visitUnion(UnionNode node, Context context) { + return visitUnorderedFanIn(node, context); + } + + private PlanNode visitUnorderedFanIn(PlanNode node, Context context) { + if (node.getChildren().size() <= 1) { + return visitPlan(node, context); + } + List children = new ArrayList<>(); + for (PlanNode child : node.getChildren()) { + children.add(child.accept(this, new Context())); + } + context.addCannotEliminateSort(true); + return node.replaceChildren(children); + } + @Override public PlanNode visitProject(ProjectNode node, Context context) { Context newContext = new Context(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java index ce4910424812..a5573159a94e 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManagerTest.java @@ -40,9 +40,11 @@ import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import org.apache.iotdb.mpp.rpc.thrift.TNewDataBlockEvent; +import org.apache.tsfile.read.common.block.TsBlock; import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; @@ -180,6 +182,13 @@ public void testCreateLocalSinkHandle() { Assert.assertEquals( ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(), ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue()); + assertProducerOwnsQueue( + spyMemoryPool, + shuffleSinkHandle, + localSourceHandle, + localFragmentInstanceId, + remoteFragmentInstanceId, + localPlanNodeId); } @Test @@ -237,5 +246,49 @@ public void testCreateLocalSourceHandle() { Assert.assertEquals( ((LocalSinkChannel) shuffleSinkHandle.getChannel(0)).getSharedTsBlockQueue(), ((LocalSourceHandle) localSourceHandle).getSharedTsBlockQueue()); + assertProducerOwnsQueue( + spyMemoryPool, + shuffleSinkHandle, + localSourceHandle, + localFragmentInstanceId, + remoteFragmentInstanceId, + localPlanNodeId); + } + + private static void assertProducerOwnsQueue( + MemoryPool pool, + ISinkHandle sink, + ISourceHandle source, + TFragmentInstanceId producer, + TFragmentInstanceId consumer, + String planNodeId) { + SharedTsBlockQueue queue = ((LocalSourceHandle) source).getSharedTsBlockQueue(); + sink.setMaxBytesCanReserve(512); + source.setMaxBytesCanReserve(128); + Assert.assertEquals(512, queue.getMaxBytesCanReserve()); + String producerId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(producer); + String consumerId = + FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(consumer); + TsBlock block = Utils.createMockTsBlock(256); + synchronized (queue) { + Assert.assertTrue(queue.add(block).isDone()); + Assert.assertEquals(256, pool.getReservedBytes()); + Assert.assertEquals( + 256, pool.getPlanNodeMemoryReservedBytes(producer.queryId, producerId, planNodeId)); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(consumer.queryId, consumerId, MemorySide.OUTPUT)); + Assert.assertSame(block, queue.remove()); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(producer.queryId, producerId, MemorySide.OUTPUT)); + } } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 2d9e5683cfdf..3cfebb1ba295 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -66,9 +66,10 @@ public class SourceHandleTest { public static void beforeClass() { maxBytesPerFI = IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + // INPUT receives half of the fragment budget, allowing five buffered blocks. IoTDBDescriptor.getInstance() .getMemoryConfig() - .setMaxBytesPerFragmentInstance(5 * MOCK_TSBLOCK_SIZE); + .setMaxBytesPerFragmentInstance(10 * MOCK_TSBLOCK_SIZE); } @AfterClass diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 5f2285e4ca91..19fc0f17d336 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -50,7 +50,8 @@ public class MemoryPoolTest { public void before() { MemoryManager memoryManager = Mockito.spy(new MemoryManager(1024L)); pool = new MemoryPool("test", memoryManager, 512L); - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } @Test @@ -87,7 +88,7 @@ public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws Exceptio // Replace the query map while the delayed cleanup still refers to the old one. pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "result", PLAN_NODE_ID, MemorySide.INPUT); Assert.assertTrue( pool.tryReserveForTest(QUERY_ID, "result", PLAN_NODE_ID, 256L, Long.MAX_VALUE)); } @@ -123,7 +124,7 @@ public void testDeregisterMissingFragmentPreservesReservedMemory() { @Test public void testDeregisterPreservesOtherFragments() { - pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID, MemorySide.INPUT); Assert.assertTrue( pool.tryReserveForTest(QUERY_ID, "other", PLAN_NODE_ID, 256L, Long.MAX_VALUE)); @@ -518,4 +519,112 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { Assert.assertTrue(r.getFuture().isDone()); Assert.assertEquals(256L, pool.getReservedBytes()); } + + @Test + public void testSharedInputBudgetAndIndependentOutput() { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 256, false) + .isReserveSuccess()); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256, false); + Assert.assertFalse(pending.isReserveSuccess()); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertEquals( + 200, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2")); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 256, false) + .isReserveSuccess()); + Assert.assertEquals(456, pool.getReservedBytes()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200); + Assert.assertTrue(pending.getFuture().isDone()); + Assert.assertEquals( + 100, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + } + + @Test + public void testSideRegistrationCannotChange() { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); + try { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.OUTPUT); + Assert.fail("conflicting side must be rejected"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + } + } + + @Test + public void testOddBudgetAndSingleRequestBoundary() { + Assert.assertEquals(256, MemorySide.INPUT.getMaxBytes(513)); + Assert.assertEquals(257, MemorySide.OUTPUT.getMaxBytes(513)); + try { + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false); + Assert.fail("request exceeds the input budget"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + } + } + + @Test + public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newFixedThreadPool(16); + java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.CountDownLatch attempted = new java.util.concurrent.CountDownLatch(16); + java.util.concurrent.CountDownLatch release = new java.util.concurrent.CountDownLatch(1); + java.util.List> tasks = new java.util.ArrayList<>(); + try { + for (int i = 0; i < 16; i++) { + String id = "pipeline-" + i; + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + start.await(); + boolean reserved = + pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128); + attempted.countDown(); + release.await(); + if (reserved) pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32); + return null; + })); + } + start.countDown(); + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + long details = 0; + for (int i = 0; i < 16; i++) { + details += + pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "pipeline-" + i); + } + Assert.assertTrue(details > 0 && details <= 128); + Assert.assertEquals( + details, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals(details, pool.getReservedBytes()); + release.countDown(); + for (java.util.concurrent.Future task : tasks) task.get(10, TimeUnit.SECONDS); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + } finally { + start.countDown(); + release.countDown(); + executor.shutdownNow(); + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java new file mode 100644 index 000000000000..68bff035fe2c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -0,0 +1,78 @@ +/* + * 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.planner; + +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; +import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; +import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; +import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; + +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class LocalExecutionPlanContextBudgetTest { + @Test + public void testFixedBudgetsInheritedByPipelinesAndWiredBySide() { + long previous = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + try { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(513); + FragmentInstanceContext instance = mock(FragmentInstanceContext.class); + when(instance.getId()) + .thenReturn(new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0")); + LocalExecutionPlanContext context = + new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); + context.setExchangeSumNum(500); + assertEquals(256, context.getInputMaxBytes()); + assertEquals(257, context.getOutputMaxBytes()); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); + LocalExecutionPlanContext child = context.createSubContext(); + assertEquals(256, child.getInputMaxBytes()); + assertEquals(257, child.getOutputMaxBytes()); + ISink sink = mock(ISink.class); + child.setISink(sink); + verify(sink).setMaxBytesCanReserve(257); + ISourceHandle remote = mock(ISourceHandle.class); + PipelineSourceHandle pipeline = mock(PipelineSourceHandle.class); + ExchangeOperator remoteOperator = mock(ExchangeOperator.class); + ExchangeOperator pipelineOperator = mock(ExchangeOperator.class); + when(remoteOperator.getSourceHandle()).thenReturn(remote); + when(pipelineOperator.getSourceHandle()).thenReturn(pipeline); + context.addExchangeOperator(remoteOperator); + child.addExchangeOperator(pipelineOperator); + context.setExchangeMemoryLimits(); + verify(remote).setMaxBytesCanReserve(256); + verify(pipeline).setMaxBytesCanReserve(257); + } finally { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previous); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java new file mode 100644 index 000000000000..3021363139c7 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java @@ -0,0 +1,149 @@ +/* + * 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.relational.planner; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; + +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class CollectInputPropertyTest { + private static PlanNode plan(String sql) { + PlanTester tester = new PlanTester(); + tester.createPlan(sql); + return tester.getFragmentPlan(0); + } + + private static List nodes(PlanNode node, Class type) { + List result = new ArrayList<>(); + if (type.isInstance(node)) result.add(type.cast(node)); + for (PlanNode child : node.getChildren()) result.addAll(nodes(child, type)); + return result; + } + + @Test + public void testAggregationAboveGlobalWindowHashesInterleavedGroups() { + for (String aggregate : new String[] {"SUM(s1)", "COUNT(DISTINCT s1)"}) { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, " + + aggregate + + " FROM (" + + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + for (AggregationNode aggregation : aggregations) { + if (!nodes(aggregation.getChild(), CollectNode.class).isEmpty()) { + assertTrue(aggregation.getPreGroupedSymbols().isEmpty()); + assertFalse(aggregation.isStreamable()); + } + } + assertTrue( + "Hash aggregation should not add a grouping sort", nodes(root, SortNode.class).isEmpty()); + } + } + + @Test + public void testPartialPreGroupingAlsoFallsBackToHash() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) FROM (" + + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 WHERE tag1 != 'shenzhen') t WHERE n > 0 GROUP BY tag1, tag2, tag3, s1"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertTrue(aggregations.get(0).getPreGroupedSymbols().isEmpty()); + } + + @Test + public void testNestedWindowRestoresPartitionContiguityWithSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, SUM(s1) OVER (" + + "PARTITION BY tag1, tag2, tag3 ORDER BY s1) AS running_sum " + + "FROM (SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1) t WHERE n > 0"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertTrue(nodes(root, GroupNode.class).isEmpty()); + List windows = nodes(root, WindowNode.class); + assertEquals(2, windows.size()); + WindowNode outer = + windows.stream() + .filter(w -> !w.getSpecification().getPartitionBy().isEmpty()) + .findFirst() + .get(); + assertFalse(nodes(outer.getChild(), SortNode.class).isEmpty()); + } + + @Test + public void testStreamSortAfterGlobalWindowUsesFullSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n " + + "FROM table1 ORDER BY tag1, tag2, tag3, s1"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertFalse(nodes(root, SortNode.class).isEmpty()); + assertTrue(nodes(root, StreamSortNode.class).isEmpty()); + } + + @Test + public void testConstantFillDoesNotRemoveRestoredSort() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 " + + "FILL METHOD CONSTANT 0 ORDER BY tag1, tag2, tag3"); + assertFalse(nodes(root, CollectNode.class).isEmpty()); + assertFalse(nodes(root, SortNode.class).isEmpty()); + assertTrue(nodes(root, StreamSortNode.class).isEmpty()); + } + + @Test + public void testUnionFinalAggregationDoesNotInheritBranchGrouping() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, SUM(s1) FROM (" + + "SELECT tag1, tag2, tag3, s1 FROM table1 UNION ALL " + + "SELECT tag1, tag2, tag3, s1 FROM table2) u GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + boolean foundFinalCollect = false; + for (AggregationNode aggregation : aggregations) { + if (aggregation.getChild() instanceof CollectNode) { + foundFinalCollect = true; + assertTrue(aggregation.getPreGroupedSymbols().isEmpty()); + } + } + assertTrue(foundFinalCollect); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java index e2326688112b..b5b69f7663f0 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/ParallelizeGroupingTest.java @@ -26,6 +26,9 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.OrderingScheme; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.ValuesNode; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.warnings.WarningCollector; @@ -36,6 +39,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExternalTsFileScanNode; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableListMultimap; import com.google.common.collect.ImmutableMap; import org.junit.Test; @@ -108,4 +112,54 @@ region, new ColumnSchema("region", STRING, false, TAG), resource.closeByQueryExecution(); } } + + @Test + public void testUnionDoesNotInheritFirstBranchGrouping() { + Symbol key = new Symbol("key"); + OrderingScheme ordering = + new OrderingScheme(ImmutableList.of(key), ImmutableMap.of(key, ASC_NULLS_LAST)); + PlanNode left = + new SortNode( + new PlanNodeId("leftSort"), + new ValuesNode(new PlanNodeId("left"), ImmutableList.of(key), ImmutableList.of()), + ordering, + false, + false); + PlanNode right = + new SortNode( + new PlanNodeId("rightSort"), + new ValuesNode(new PlanNodeId("right"), ImmutableList.of(key), ImmutableList.of()), + ordering, + false, + false); + UnionNode union = + new UnionNode( + new PlanNodeId("union"), + ImmutableList.of(left, right), + ImmutableListMultimap.of(key, key, key, key), + ImmutableList.of(key)); + assertTrue( + optimizeGroup(new GroupNode(new PlanNodeId("group"), union, ordering, 1)) + instanceof SortNode); + // An explicit sort above the fan-in restores the grouping property. + SortNode restored = new SortNode(new PlanNodeId("restored"), union, ordering, false, false); + assertTrue( + optimizeGroup(new GroupNode(new PlanNodeId("group2"), restored, ordering, 1)) + instanceof GroupNode); + } + + private PlanNode optimizeGroup(GroupNode group) { + return new ParallelizeGrouping() + .optimize( + group, + new PlanOptimizer.Context( + SESSION_INFO, + new Analysis(null, Collections.emptyMap()), + TEST_MATADATA, + new MPPQueryContext(new QueryId("group_union")), + new SymbolAllocator(), + new QueryId("group_union"), + WarningCollector.NOOP, + PlanOptimizersStatsCollector.createPlanOptimizersStatsCollector())); + } } From 0df3a3fb47206006815a34a4856532a0c77410ad Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 14 Oct 2026 14:53:56 +0800 Subject: [PATCH 03/15] Fix exchange reservation wakeup race and strengthen fan-in tests --- ...2026-10-14-exchange-fanin-budget-design.md | 22 ++- .../it/query/recent/IoTDBCollectFanInIT.java | 23 +++ .../operator/process/CollectOperatorTest.java | 102 ++++++++++++ .../execution/memory/MemoryPool.java | 51 +++--- .../execution/memory/MemoryPoolTest.java | 154 +++++++++++++++++- .../planner/CollectInputPropertyTest.java | 28 ++++ 6 files changed, 353 insertions(+), 27 deletions(-) diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index a817039cab4b..e191622940dc 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -72,7 +72,7 @@ SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除 ## 实现验证记录（2026-10-14） -第一版实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。 +上述约定范围的实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。 集成验证包含普通与 DISTINCT 聚合、Union 分阶段聚合、嵌套窗口、Fill、排序/TopK、Join 依赖、树模型以及树表视图。 构建先执行以下干净构建，生成当前分支的节点包并安装本地依赖： @@ -97,3 +97,23 @@ mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfa 使用离线模式是为了避免远端 SNAPSHOT ZIP 下载停滞；节点执行的是当前分支重新打包的代码。 Spotless、Checkstyle、新增文件许可证及中英文消息键名/占位符一致性检查通过。 + +## 自查与补测（2026-10-14） + +自查复现了 MemoryPool 的等待申请漏唤醒：一次申请额度不足后，旧申请先释放全部额度， +新申请才进入等待队列；释放线程看不到新申请，导致空闲额度下仍无限等待。 +回归测试用 latch 固定这一时序，在修复前稳定失败。 +修复在申请入队后再次检查额度，并与 free/tryCancel 共用该 future 的锁，避免重复计费。 + +新增 8 项单元测试：等待申请发布竞态、取消不占用共享额度、16 路阻塞申请反复获得额度后账目归零； +空 Collect、全部阻塞时异常唤醒、并发就绪后每路恰好消费一次； +Collect 后 TopK 不误用有序输入，以及有效分支保留流式聚合。 +增强关闭异常测试，检查其他输入仍被关闭且后续异常被保留。 +新增一项集成测试，验证全局窗口后按字段倒序的 TopK 跨设备结果及完整窗口计数。 + +补测后的 118 项定向单测通过（14 个测试类，零失败、零跳过）；全 reactor 中英文编译通过。 +在前述英文回归命令中加入 `OrderingPropertyPropagationTest,TopKRuntimeFilterDistributedPlanTest`， +同时验证有效排序优化与 TopK runtime filter 不受影响。 +重新执行 `mvn -o install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests` +后运行前述两条 verify 命令；IoTDBCollectFanInIT 的 7 项测试在单节点和三 DataNode 集群中分别全部通过。 +Spotless、Checkstyle、许可证与 `git diff --check` 检查通过。 diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java index adc72b6046e5..8667f01ddfd7 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -273,4 +273,27 @@ public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception { } } } + + @Test + public void testTopKAfterGlobalWindowConsumesUnorderedInput() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + int limit = DEVICES + 5; + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value, COUNT(*) OVER () AS n FROM readings1 " + + "ORDER BY value DESC, device LIMIT " + + limit)) { + for (int row = 0; row < limit; row++) { + assertTrue(result.next()); + assertEquals(device(row % DEVICES), result.getString(1)); + assertEquals(zone(row % DEVICES), result.getString(2)); + assertEquals(ROWS - row / DEVICES, result.getInt(3)); + assertEquals((long) DEVICES * ROWS, result.getLong(4)); + } + assertFalse(result.next()); + } + } + } } diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java index abf7a82acf3d..b446f9c999b9 100644 --- a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -36,9 +36,16 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -233,6 +240,7 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep Child a = new Child(block(1)); Child b = new Child(block(2)); a.closeFailure = new IOException("close failure"); + b.closeFailure = new IOException("second close failure"); CollectOperator collect = collect(a, b); ListenableFuture pending = collect.isBlocked(); try { @@ -240,6 +248,8 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep fail("must propagate close failure"); } catch (IOException e) { assertSame(a.closeFailure, e); + assertEquals(1, e.getSuppressed().length); + assertSame(b.closeFailure, e.getSuppressed()[0]); } pending.get(5, TimeUnit.SECONDS); assertEquals(1, a.closes); @@ -299,6 +309,98 @@ public void testMemoryIncludesAllLiveRetainedStates() throws Exception { assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); } + @Test + public void testEmptyCollectIsAlreadyFinished() throws Exception { + CollectOperator collect = collect(); + assertFalse(collect.hasNext()); + assertTrue(collect.isFinished()); + assertTrue(collect.isBlocked().isDone()); + assertNull(collect.next()); + assertEquals(0, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateMaxReturnSize()); + collect.close(); + collect.close(); + } + + @Test + public void testFailureWakesAllPendingChildren() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + CollectOperator collect = collect(a, b); + ListenableFuture pending = collect.isBlocked(); + assertFalse(pending.isDone()); + IOException expected = new IOException("blocked child failure"); + b.wait.delegate.setException(expected); + try { + pending.get(5, TimeUnit.SECONDS); + fail("child failure must wake the driver"); + } catch (ExecutionException e) { + assertSame(expected, e.getCause()); + } + try { + collect.next(); + fail("child failure must propagate without advancing other children"); + } catch (IOException e) { + assertSame(expected, e); + } finally { + collect.close(); + } + assertEquals(0, a.advances); + assertEquals(0, b.advances); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Exception { + int count = 16; + List children = new ArrayList<>(); + for (int i = 0; i < count; i++) children.add(new Child(block(i))); + CollectOperator collect = new CollectOperator(new TestContext(), new ArrayList<>(children)); + ListenableFuture pending = collect.isBlocked(); + assertFalse(pending.isDone()); + ExecutorService executor = Executors.newFixedThreadPool(count); + CountDownLatch start = new CountDownLatch(1); + List> completions = new ArrayList<>(); + try { + for (Child child : children) { + completions.add( + executor.submit( + () -> { + start.await(); + child.ready(); + return null; + })); + } + start.countDown(); + pending.get(5, TimeUnit.SECONDS); + for (Future completion : completions) completion.get(5, TimeUnit.SECONDS); + Set consumed = new HashSet<>(); + for (int round = 0; round < 2 * count && collect.hasNext(); round++) { + assertTrue(collect.isBlocked().isDone()); + TsBlock result = collect.next(); + if (result != null) { + int childId = result.getColumn(0).getInt(0); + assertTrue("child emitted more than once", consumed.add(childId)); + children.get(childId).hasNext = false; + children.get(childId).finished = true; + } + } + assertEquals(count, consumed.size()); + assertTrue(collect.isFinished()); + assertNull(collect.next()); + for (Child child : children) { + assertEquals(1, child.advances); + assertEquals(1, child.closes); + } + } finally { + start.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + collect.close(); + } + } + private static final class CountingFuture extends ForwardingListenableFuture.SimpleForwardingListenableFuture { private final SettableFuture delegate; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 56c9974764a5..07eced179d94 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -365,10 +365,15 @@ public MemoryReservationResult reserveWithPriority( DataNodeQueryMessages.BLOCKED_RESERVE_REQUEST_ARG_BYTES_MEMORY_FOR_PLANNODEIDARG, bytesToReserve, planNodeId); - ListenableFuture result = + MemoryReservationFuture result = MemoryReservationFuture.create( queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); - memoryReservationFutures.add((MemoryReservationFuture) result); + memoryReservationFutures.add(result); + // A concurrent free may have missed this request before it was published. Retry after + // publication, under the same future lock used by free and tryCancel to avoid charging twice. + if (tryCompleteReservation(result)) { + memoryReservationFutures.remove(result); + } return new MemoryReservationResult(result, false, bytesToReserve); } } @@ -463,27 +468,35 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l Iterator> iterator = memoryReservationFutures.iterator(); while (iterator.hasNext()) { MemoryReservationFuture future = iterator.next(); - synchronized (future) { - if (future.isCancelled() || future.isDone()) { - iterator.remove(); - continue; - } - long bytesToReserve = future.getBytesToReserve(); - String curQueryId = future.getQueryId(); - String curFragmentInstanceId = future.getFragmentInstanceId(); - String curPlanNodeId = future.getPlanNodeId(); - long maxBytesCanReserve = future.getMaxBytesCanReserve(); - if (tryReserve( - curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve, maxBytesCanReserve)) { - future.set(null); - iterator.remove(); - } else { - rollbackReserve(curQueryId, curFragmentInstanceId, curPlanNodeId, bytesToReserve); - } + if (tryCompleteReservation(future)) { + iterator.remove(); } } } + private boolean tryCompleteReservation(MemoryReservationFuture future) { + synchronized (future) { + if (future.isDone()) { + return true; + } + long bytesToReserve = future.getBytesToReserve(); + String queryId = future.getQueryId(); + String fragmentInstanceId = future.getFragmentInstanceId(); + String planNodeId = future.getPlanNodeId(); + if (tryReserve( + queryId, + fragmentInstanceId, + planNodeId, + bytesToReserve, + future.getMaxBytesCanReserve())) { + future.set(null); + return true; + } + rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); + return false; + } + } + public long getQueryMemoryReservedBytes(String queryId) { Map fragments = queryMemoryReservations.get(queryId); if (fragments == null) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 19fc0f17d336..0254fa1ac33c 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -32,8 +32,15 @@ import java.lang.management.ManagementFactory; import java.lang.management.ThreadInfo; import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.FutureTask; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import static org.awaitility.Awaitility.await; @@ -582,12 +589,11 @@ public void testOddBudgetAndSingleRequestBoundary() { @Test public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { - java.util.concurrent.ExecutorService executor = - java.util.concurrent.Executors.newFixedThreadPool(16); - java.util.concurrent.CountDownLatch start = new java.util.concurrent.CountDownLatch(1); - java.util.concurrent.CountDownLatch attempted = new java.util.concurrent.CountDownLatch(16); - java.util.concurrent.CountDownLatch release = new java.util.concurrent.CountDownLatch(1); - java.util.List> tasks = new java.util.ArrayList<>(); + ExecutorService executor = Executors.newFixedThreadPool(16); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(16); + CountDownLatch release = new CountDownLatch(1); + List> tasks = new ArrayList<>(); try { for (int i = 0; i < 16; i++) { String id = "pipeline-" + i; @@ -618,7 +624,7 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); Assert.assertEquals(details, pool.getReservedBytes()); release.countDown(); - for (java.util.concurrent.Future task : tasks) task.get(10, TimeUnit.SECONDS); + for (Future task : tasks) task.get(10, TimeUnit.SECONDS); Assert.assertEquals(0, pool.getReservedBytes()); Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); } finally { @@ -627,4 +633,138 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { executor.shutdownNow(); } } + + @Test + public void testReleaseBeforePendingReservationIsPublished() throws Exception { + pool = Mockito.spy(pool); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false) + .isReserveSuccess()); + CountDownLatch failedAttempt = new CountDownLatch(1); + CountDownLatch released = new CountDownLatch(1); + AtomicBoolean firstAttempt = new AtomicBoolean(true); + Mockito.doAnswer( + invocation -> { + boolean granted = (boolean) invocation.callRealMethod(); + if (!granted && firstAttempt.compareAndSet(true, false)) { + failedAttempt.countDown(); + Assert.assertTrue(released.await(10, TimeUnit.SECONDS)); + } + return granted; + }) + .when(pool) + .tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256); + FutureTask reservation = + new FutureTask<>( + () -> + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256, 256, false)); + Thread reserver = new Thread(reservation, "reserve-shared-input"); + reserver.start(); + try { + Assert.assertTrue(failedAttempt.await(10, TimeUnit.SECONDS)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + released.countDown(); + MemoryReservationResult result = reservation.get(10, TimeUnit.SECONDS); + Assert.assertTrue( + "Releasing the last owner must wake the pending request", result.getFuture().isDone()); + result.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertEquals(256, pool.getReservedBytes()); + Assert.assertEquals( + 256, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 256, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2")); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } finally { + released.countDown(); + reserver.join(10000); + Assert.assertFalse(reserver.isAlive()); + } + } + + @Test + public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception { + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false); + MemoryReservationResult cancelled = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", 128, 256, false); + MemoryReservationResult waiting = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256, 256, false); + Assert.assertFalse(cancelled.getFuture().isDone()); + Assert.assertFalse(waiting.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128, 256, false) + .isReserveSuccess()); + Assert.assertEquals(128, pool.tryCancel(cancelled.getFuture())); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + waiting.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertTrue(cancelled.getFuture().isCancelled()); + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled")); + Assert.assertEquals( + 256, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals( + 128, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.OUTPUT)); + Assert.assertEquals(384, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting", 256); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 128); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception { + int count = 16; + ExecutorService executor = Executors.newFixedThreadPool(count); + CountDownLatch attempted = new CountDownLatch(count); + List> tasks = new ArrayList<>(); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128, 128, false); + try { + for (int i = 0; i < count; i++) { + String id = "waiting-" + i; + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + for (int round = 0; round < 50; round++) { + MemoryReservationResult result = + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128, false); + if (round == 0) attempted.countDown(); + result.getFuture().get(10, TimeUnit.SECONDS); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, result.getReservedBytes()); + } + return null; + })); + } + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + Assert.assertEquals(count, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 128); + for (Future task : tasks) task.get(30, TimeUnit.SECONDS); + for (int i = 0; i < count; i++) { + Assert.assertEquals( + 0, pool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "waiting-" + i)); + } + Assert.assertEquals( + 0, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } finally { + executor.shutdownNow(); + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java index 3021363139c7..5bae986bf476 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; import org.junit.Test; @@ -146,4 +147,31 @@ public void testUnionFinalAggregationDoesNotInheritBranchGrouping() { } assertTrue(foundFinalCollect); } + + @Test + public void testTopKAboveGlobalWindowDoesNotAssumeOrderedInput() { + PlanNode root = + plan( + "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 " + + "ORDER BY s1 DESC, tag1, tag2, tag3 LIMIT 5"); + List topKs = nodes(root, TopKNode.class); + assertFalse(topKs.isEmpty()); + boolean foundUnorderedInput = false; + for (TopKNode topK : topKs) { + if (!nodes(topK.getChildren().get(0), CollectNode.class).isEmpty()) { + foundUnorderedInput = true; + assertFalse(topK.isChildrenDataInOrder()); + } + } + assertTrue(foundUnorderedInput); + } + + @Test + public void testOrderedScanBranchesKeepStreamingAggregation() { + PlanNode root = + plan("SELECT tag1, tag2, tag3, SUM(s1 + 1) FROM table1 GROUP BY tag1, tag2, tag3"); + List aggregations = nodes(root, AggregationNode.class); + assertFalse(aggregations.isEmpty()); + assertTrue(aggregations.stream().anyMatch(AggregationNode::isStreamable)); + } } From 75cd4e7f5afd649fbf307796ed7b0ea75b65e4f8 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 08:15:10 +0800 Subject: [PATCH 04/15] Fix unordered result assertions for ready Collect inputs Validate unordered results and EXPLAIN semantics independently on every DataNode while retaining ordered assertions. Add duplicate-preserving checks and deterministic regressions, and resolve receiver column names from metadata. --- .../itbase/runtime/ClusterTestResultSet.java | 12 ++ .../apache/iotdb/db/it/utils/TestUtils.java | 46 ++++-- .../it/triple/IoTDBShowReceiversTripleIT.java | 84 ++++++---- .../it/db/it/IoTDBMultiDeviceTableIT.java | 27 ++-- .../udf/IoTDBUserDefinedTableFunctionIT.java | 45 +++--- .../it/insertquery/IoTDBInsertQueryIT.java | 19 ++- .../query/recent/IoTDBClusterResultSetIT.java | 152 ++++++++++++++++++ .../it/query/recent/IoTDBCollectFanInIT.java | 25 +++ .../it/query/recent/IoTDBCteIT.java | 82 +++++----- .../query/recent/IoTDBTableAggregationIT.java | 28 ++-- .../IoTDBTopKRuntimeFilterScanPathIT.java | 29 ++-- 11 files changed, 400 insertions(+), 149 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java diff --git a/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java b/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java index 98686d436be1..a9823b28572a 100644 --- a/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java +++ b/integration-test/src/main/java/org/apache/iotdb/itbase/runtime/ClusterTestResultSet.java @@ -40,6 +40,7 @@ import java.sql.Time; import java.sql.Timestamp; import java.util.Calendar; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -68,6 +69,17 @@ public ClusterTestResultSet( this.env = env; } + /** + * Returns each node's result for assertions that do not require identical row order or plan text. + * Every returned result must be checked. The caller must close the original result set and must + * not mix per-node iteration with iteration through the cluster wrapper. + */ + public static List getNodeResultSets(ResultSet resultSet) { + return resultSet instanceof ClusterTestResultSet + ? Collections.unmodifiableList(((ClusterTestResultSet) resultSet).resultSets) + : Collections.singletonList(resultSet); + } + @Override public boolean next() throws SQLException { RequestDelegate delegate = createRequestDelegate(); diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java index bf1cb4376bb9..a85aacc4c368 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java @@ -63,6 +63,7 @@ import static org.apache.iotdb.itbase.constant.TestConstant.TIMESTAMP_STR; import static org.apache.iotdb.itbase.env.BaseEnv.TABLE_SQL_DIALECT; import static org.apache.iotdb.itbase.env.BaseEnv.TREE_SQL_DIALECT; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.awaitility.Awaitility.await; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -753,6 +754,15 @@ public static void assertResultSetEqual( final String expectedHeader, final Collection expectedResult, final boolean printResult) { + if (expectedResult instanceof Set) { + List nodeResults = getNodeResultSets(actualResultSet); + if (nodeResults.size() != 1 || nodeResults.get(0) != actualResultSet) { + for (ResultSet nodeResult : nodeResults) { + assertResultSetEqual(nodeResult, expectedHeader, expectedResult, printResult); + } + return; + } + } try { final ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData(); final StringBuilder header = new StringBuilder(); @@ -814,27 +824,31 @@ public static void assertResultSetEqual( String expectedHeader, Set expectedRetSet, Consumer consumer) { - try { - ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData(); + assertResultSetEqual(actualResultSet, expectedHeader, expectedRetSet, false); + } + + /** Checks unordered rows on every node, preserving duplicate multiplicity. */ + public static void assertResultSetEqualUnordered( + ResultSet resultSet, String expectedHeader, List expectedRows) throws SQLException { + List sortedExpected = new ArrayList<>(expectedRows); + Collections.sort(sortedExpected); + for (ResultSet nodeResult : getNodeResultSets(resultSet)) { + ResultSetMetaData metadata = nodeResult.getMetaData(); StringBuilder header = new StringBuilder(); - for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) { - header.append(resultSetMetaData.getColumnName(i)).append(","); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + header.append(metadata.getColumnName(i)).append(","); } assertEquals(expectedHeader, header.toString()); - - Set actualRetSet = new HashSet<>(); - - while (actualResultSet.next()) { - StringBuilder builder = new StringBuilder(); - for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) { - builder.append(actualResultSet.getString(i)).append(","); + List actualRows = new ArrayList<>(); + while (nodeResult.next()) { + StringBuilder row = new StringBuilder(); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + row.append(nodeResult.getString(i)).append(","); } - actualRetSet.add(builder.toString()); + actualRows.add(row.toString()); } - assertStringSetEqual(expectedRetSet, actualRetSet); - } catch (Exception e) { - e.printStackTrace(); - Assert.fail(String.valueOf(e)); + Collections.sort(actualRows); + assertEquals(sortedExpected, actualRows); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java index 0712387775b9..fda9935e3191 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java @@ -44,6 +44,7 @@ import java.sql.Connection; import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; @@ -54,6 +55,8 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; + @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT3.class}) public class IoTDBShowReceiversTripleIT { @@ -214,39 +217,44 @@ private boolean hasReceiversOnDifferentDataNodes( final String secondPipeName, final int secondReceiverDataNodeId) throws SQLException { - boolean hasFirstPipe = false; - boolean hasSecondPipe = false; - final Set receiverDataNodeIds = new HashSet<>(); try (final Connection connection = receiverEnv.getConnection( SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); final Statement statement = connection.createStatement(); - final ResultSet resultSet = statement.executeQuery(sql)) { - while (resultSet.next()) { - final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); - if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) - || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) - || pipeIds == null) { - continue; - } - final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); - if (receiverDataNodeId == firstReceiverDataNodeId - && pipeIds.contains(firstPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasFirstPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); + final ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + boolean hasFirstPipe = false; + boolean hasSecondPipe = false; + final Set receiverDataNodeIds = new HashSet<>(); + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) + || pipeIds == null) { + continue; + } + final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); + if (receiverDataNodeId == firstReceiverDataNodeId + && pipeIds.contains(firstPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasFirstPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } + if (receiverDataNodeId == secondReceiverDataNodeId + && pipeIds.contains(secondPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasSecondPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } } - if (receiverDataNodeId == secondReceiverDataNodeId - && pipeIds.contains(secondPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasSecondPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); + if (!hasFirstPipe || !hasSecondPipe || receiverDataNodeIds.size() < 2) { + return false; } } } - return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2; + return true; } private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception { @@ -268,20 +276,28 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception private static String getString( final ResultSet resultSet, final String treeColumnName, final String tableColumnName) throws SQLException { - try { - return resultSet.getString(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getString(tableColumnName); - } + return resultSet.getString(getColumnIndex(resultSet, treeColumnName, tableColumnName)); } private static int getInt( final ResultSet resultSet, final String treeColumnName, final String tableColumnName) throws SQLException { - try { - return resultSet.getInt(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getInt(tableColumnName); + return resultSet.getInt(getColumnIndex(resultSet, treeColumnName, tableColumnName)); + } + + private static int getColumnIndex( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + // SHOW RECEIVERS and information_schema.receivers expose different column names. + // Resolve from metadata because JDBC getters do not all use SQLException for unknown names. + final ResultSetMetaData metadata = resultSet.getMetaData(); + for (int column = 1; column <= metadata.getColumnCount(); column++) { + final String name = metadata.getColumnName(column); + if (treeColumnName.equals(name) || tableColumnName.equals(name)) { + return column; + } } + throw new AssertionError( + "Missing receiver column: " + treeColumnName + " / " + tableColumnName); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java index 16d92c213436..096c1383c10d 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java @@ -37,6 +37,7 @@ import java.util.HashMap; import java.util.Map; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -269,20 +270,22 @@ private void testSelectAll() { try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); Statement statement = connection.createStatement()) { statement.execute("use \"test\""); - Map lastTimeMap = new HashMap<>(); - try (ResultSet resultSet = statement.executeQuery(selectSql)) { - int cnt = 0; - while (resultSet.next()) { - String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2"); - long before = lastTimeMap.getOrDefault(tag, -1L); - long cur = resultSet.getTimestamp("time").getTime(); - if (cur <= before) { - fail("time order wrong!"); + try (ResultSet clusterResult = statement.executeQuery(selectSql)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + Map lastTimeMap = new HashMap<>(); + int cnt = 0; + while (resultSet.next()) { + String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2"); + long before = lastTimeMap.getOrDefault(tag, -1L); + long cur = resultSet.getTimestamp("time").getTime(); + if (cur <= before) { + fail("time order wrong!"); + } + lastTimeMap.put(tag, cur); + cnt++; } - lastTimeMap.put(tag, cur); - cnt++; + assertEquals(16030, cnt); } - assertEquals(16030, cnt); } } catch (Exception e) { e.printStackTrace(); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java index 9db3acb13e09..25967b8caeb8 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java @@ -41,6 +41,7 @@ import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.fail; @RunWith(IoTDBTestRunner.class) @@ -200,37 +201,39 @@ public void testLargeResultIsSplitWithoutDataLoss() throws Exception { "large_result", "org.apache.iotdb.db.query.udf.example.relational.LargeResultTableFunction"); - Set returnedRows = new HashSet<>(); try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); - try (ResultSet resultSet = + try (ResultSet clusterResult = statement.executeQuery( "SELECT * FROM large_result(vehicle, " + LARGE_RESULT_REPEAT_COUNT + ", " + LARGE_RESULT_PAYLOAD_SIZE + ")")) { - while (resultSet.next()) { - int repeatIndex = resultSet.getInt("repeat_index"); - long time = resultSet.getLong("time"); - Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT); - Assert.assertEquals( - repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE), - resultSet.getString("payload")); - assertPassThroughColumns(resultSet, time); - Assert.assertTrue( - "Duplicate result row for time " + time + " and repeat index " + repeatIndex, - returnedRows.add(time + ":" + repeatIndex)); - } - } - } + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + Set returnedRows = new HashSet<>(); + while (resultSet.next()) { + int repeatIndex = resultSet.getInt("repeat_index"); + long time = resultSet.getLong("time"); + Assert.assertTrue(repeatIndex >= 0 && repeatIndex < LARGE_RESULT_REPEAT_COUNT); + Assert.assertEquals( + repeatIndex + ":" + "x".repeat(LARGE_RESULT_PAYLOAD_SIZE), + resultSet.getString("payload")); + assertPassThroughColumns(resultSet, time); + Assert.assertTrue( + "Duplicate result row for time " + time + " and repeat index " + repeatIndex, + returnedRows.add(time + ":" + repeatIndex)); + } - long[] inputTimes = new long[] {1, 2, 3, 5}; - Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size()); - for (long time : inputTimes) { - for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) { - Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex)); + long[] inputTimes = new long[] {1, 2, 3, 5}; + Assert.assertEquals(inputTimes.length * LARGE_RESULT_REPEAT_COUNT, returnedRows.size()); + for (long time : inputTimes) { + for (int repeatIndex = 0; repeatIndex < LARGE_RESULT_REPEAT_COUNT; repeatIndex++) { + Assert.assertTrue(returnedRows.contains(time + ":" + repeatIndex)); + } + } + } } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java index 174b2d4711ae..e56aadf5c86a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java @@ -48,6 +48,7 @@ import java.util.List; import java.util.Locale; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; @@ -279,13 +280,17 @@ public void testPartialColumn() throws SQLException { statement.execute("DELETE FROM vehicle1"); statement.execute( "INSERT INTO vehicle1(time, s2, manufacturer) SELECT time, s2, manufacturer FROM vehicle0"); - resultSet = statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1"); - while (resultSet.next()) { - for (int i = 1; i <= 4; i++) { - if (i % 2 == 1) { - assertNull(resultSet.getString(i)); - } else { - assertNotNull(resultSet.getString(i)); + try (ResultSet clusterResult = + statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1")) { + for (ResultSet nodeResult : getNodeResultSets(clusterResult)) { + while (nodeResult.next()) { + for (int i = 1; i <= 4; i++) { + if (i % 2 == 1) { + assertNull(nodeResult.getString(i)); + } else { + assertNotNull(nodeResult.getString(i)); + } + } } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java new file mode 100644 index 000000000000..a29c9793b72d --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java @@ -0,0 +1,152 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.db.it.utils.TestUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.itbase.runtime.ClusterTestResultSet; + +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import javax.sql.rowset.CachedRowSet; +import javax.sql.rowset.RowSetMetaDataImpl; +import javax.sql.rowset.RowSetProvider; + +import java.lang.reflect.Proxy; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Types; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; + +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +/** Deterministic coverage of cross-node assertions without depending on network arrival order. */ +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBClusterResultSetIT { + + @Test + public void testUnorderedSetChecksEveryNode() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,")); + } + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "c"})) { + assertThrows( + AssertionError.class, + () -> TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"))); + } + } + + @Test + public void testUnorderedSetConsumerOverloadChecksEveryNode() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"), message -> {}); + } + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"a"})) { + assertThrows( + AssertionError.class, + () -> + TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"), message -> {})); + } + } + + @Test + public void testUnorderedRowsPreserveDuplicateMultiplicity() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "a", "b"}, new String[] {"b", "a", "a"})) { + TestUtils.assertResultSetEqualUnordered(result, "value,", Arrays.asList("a,", "a,", "b,")); + } + for (String[] rows : new String[][] {{"a", "b"}, {"a", "a", "a", "b"}}) { + try (ResultSet result = cluster(new String[] {"a", "a", "b"}, rows)) { + assertThrows( + AssertionError.class, + () -> + TestUtils.assertResultSetEqualUnordered( + result, "value,", Arrays.asList("a,", "a,", "b,"))); + } + } + } + + @Test + public void testOrderedRowsStillRejectDifferentNodeOrder() throws SQLException { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + assertThrows( + AssertionError.class, + () -> TestUtils.assertResultSetEqual(result, "value,", Arrays.asList("a,", "b,"))); + } + } + + @Test + public void testSingleNodeAndNullValues() throws SQLException { + try (ResultSet result = rows("a", null, "a")) { + assertEquals(1, getNodeResultSets(result).size()); + TestUtils.assertResultSetEqualUnordered(result, "value,", Arrays.asList("null,", "a,", "a,")); + } + try (ResultSet result = cluster(new String[] {"a"})) { + assertEquals(1, getNodeResultSets(result).size()); + TestUtils.assertResultSetEqual(result, "value,", Set.of("a,")); + } + } + + private static ResultSet cluster(String[]... nodeRows) throws SQLException { + List statements = new ArrayList<>(); + List endpoints = new ArrayList<>(); + for (String[] values : nodeRows) { + ResultSet result = rows(values); + Statement statement = + (Statement) + Proxy.newProxyInstance( + Statement.class.getClassLoader(), + new Class[] {Statement.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("executeQuery")) { + return result; + } + throw new UnsupportedOperationException(method.getName()); + }); + statements.add(statement); + endpoints.add("node-" + endpoints.size()); + } + return new ClusterTestResultSet(statements, endpoints, "SELECT value", 10, null); + } + + private static CachedRowSet rows(String... values) throws SQLException { + CachedRowSet result = RowSetProvider.newFactory().createCachedRowSet(); + RowSetMetaDataImpl metadata = new RowSetMetaDataImpl(); + metadata.setColumnCount(1); + metadata.setColumnName(1, "value"); + metadata.setColumnType(1, Types.VARCHAR); + result.setMetaData(metadata); + for (String value : values) { + result.afterLast(); + result.moveToInsertRow(); + result.updateString(1, value); + result.insertRow(); + result.moveToCurrentRow(); + } + result.beforeFirst(); + return result; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java index 8667f01ddfd7..f9de8a0664fb 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -33,7 +33,10 @@ import java.sql.Connection; import java.sql.ResultSet; import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import static org.apache.iotdb.db.it.utils.TestUtils.assertResultSetEqualUnordered; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -83,6 +86,28 @@ public static void setUp() throws Exception { } } + @Test + public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Exception { + List expectedRows = new ArrayList<>(); + for (int device = 0; device < DEVICES; device++) { + for (int value = 1; value <= ROWS; value++) { + String row = device(device) + "," + zone(device) + "," + value + ","; + expectedRows.add(row); + expectedRows.add(row); + } + } + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery( + "SELECT device, zone, value FROM readings1 UNION ALL " + + "SELECT device, zone, value FROM readings1")) { + assertResultSetEqualUnordered(result, "device,zone,value,", expectedRows); + } + } + } + @Test public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception { try (Connection connection = EnvFactory.getEnv().getConnection(); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java index 8c0b52acc944..87571cfb05db 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java @@ -59,6 +59,7 @@ import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -362,19 +363,22 @@ public void testNestExplain1() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertFalse(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertFalse(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } @@ -389,19 +393,22 @@ public void testNestExplain2() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertTrue(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertTrue(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } @@ -416,19 +423,22 @@ public void testExplainJoin() throws SQLException { statement.execute("USE testdb"); // explain - ResultSet resultSet = statement.executeQuery(sql); - ResultSetMetaData metaData = resultSet.getMetaData(); - assertEquals(metaData.getColumnCount(), 1); - assertEquals(metaData.getColumnName(1), "distribution plan"); - - StringBuilder sb = new StringBuilder(); - while (resultSet.next()) { - sb.append(resultSet.getString(1)).append(System.lineSeparator()); + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + assertEquals(metaData.getColumnCount(), 1); + assertEquals(metaData.getColumnName(1), "distribution plan"); + + StringBuilder sb = new StringBuilder(); + while (resultSet.next()) { + sb.append(resultSet.getString(1)).append(System.lineSeparator()); + } + String result = sb.toString(); + assertFalse(result.contains("CTE Query : 'cte1'")); + assertTrue(result.contains("CTE Query : 'cte2'")); + assertTrue(result.contains("Main Query")); + } } - String result = sb.toString(); - assertFalse(result.contains("CTE Query : 'cte1'")); - assertTrue(result.contains("CTE Query : 'cte2'")); - assertTrue(result.contains("Main Query")); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java index 518b3da1fdd2..99188903b99a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java @@ -44,6 +44,7 @@ import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; +import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.apache.iotdb.relational.it.db.it.IoTDBMultiTAGsWithAttributesTableIT.buildHeaders; import static org.apache.iotdb.relational.it.db.it.IoTDBMultiTAGsWithAttributesTableIT.repeatTest; @@ -5902,19 +5903,22 @@ private void assertRateQueryUsesIntermediateAggregation(String query) { try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); - try (ResultSet resultSet = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) { - Assert.assertTrue(resultSet.next()); - String plan = resultSet.getString(1); - // The PARTIAL/FINAL split only exists when the scanned data spans multiple DataRegions. - // Under randomized partition allocation (e.g. the SHUFFLE strategy used by - // IoTDBTableAggregation2IT), all devices may occasionally land in one single region, - // where a SINGLE-step aggregation is the correct plan, so the check must be skipped. - if (countDistinctRegionIds(plan) < 2) { - return; + try (ResultSet clusterResult = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) { + for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + Assert.assertTrue(resultSet.next()); + String plan = resultSet.getString(1); + // The PARTIAL/FINAL split only exists when the scanned data spans multiple DataRegions. + // Under randomized partition allocation (e.g. the SHUFFLE strategy used by + // IoTDBTableAggregation2IT), all devices may occasionally land in one single region, + // where a SINGLE-step aggregation is the correct plan, so the check must be skipped. + if (countDistinctRegionIds(plan) < 2) { + continue; + } + Assert.assertTrue( + "Expected a PARTIAL aggregation in plan: " + plan, plan.contains("PARTIAL")); + Assert.assertTrue( + "Expected a FINAL aggregation in plan: " + plan, plan.contains("FINAL")); } - Assert.assertTrue( - "Expected a PARTIAL aggregation in plan: " + plan, plan.contains("PARTIAL")); - Assert.assertTrue("Expected a FINAL aggregation in plan: " + plan, plan.contains("FINAL")); } } catch (Exception exception) { throw new AssertionError(exception); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java index 7697bc81649b..a0a638dcfe0e 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java @@ -41,11 +41,14 @@ import java.io.IOException; import java.nio.file.Files; import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import static org.apache.iotdb.db.it.utils.TestUtils.assertResultSetEqualUnordered; import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; @@ -131,19 +134,23 @@ public void orderByTimeAscLimitOnTable() { } @Test - public void unionAllSiblingTopKOrderByTimeLimit() { - String[] expectedHeader = new String[] {"time", "device_id", "s1"}; - String[] retArray = - new String[] { - "1970-01-01T00:00:00.030Z,d1,111,", "1970-01-01T00:00:00.032Z,d3,333,", - }; - tableResultSetEqualTest( + public void unionAllSiblingTopKOrderByTimeLimit() throws SQLException { + // Each branch has an ORDER BY, but UNION ALL does not order the combined result. + String sql = "(SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd1' ORDER BY time DESC LIMIT 1)" + " UNION ALL" - + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)", - expectedHeader, - retArray, - TABLE_DATABASE); + + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)"; + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + connection.setClientInfo("time_zone", "+00:00"); + statement.execute("USE " + TABLE_DATABASE); + try (ResultSet resultSet = statement.executeQuery(sql)) { + assertResultSetEqualUnordered( + resultSet, + "time,device_id,s1,", + Arrays.asList("1970-01-01T00:00:00.030Z,d1,111,", "1970-01-01T00:00:00.032Z,d3,333,")); + } + } } @Test From 4ed7b47b33f601a2f50e8a235b3c0c51d55df854 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 08:15:11 +0800 Subject: [PATCH 05/15] Fix occupied-port detection and isolate client manager test ports Read the address column rather than the socket state from lsof output and cover a real listening socket. Allow ClientManagerTest to use an explicitly configured port when local services occupy its default. --- .../apache/iotdb/it/env/cluster/EnvUtils.java | 3 +- .../it/query/recent/IoTDBTestEnvPortIT.java | 47 +++++++++++++++++++ .../commons/client/ClientManagerTest.java | 3 +- 3 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java index b69e86c12bd4..e6c58e251f59 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java @@ -180,7 +180,8 @@ public static Map listPortOccupationWindows(final List p */ public static Map listPortOccupationUnix(final List ports) throws IOException { - return listPortOccupation(ports, "lsof -iTCP -sTCP:LISTEN,TIME_WAIT -P -n", 10, 9, 1); + // The last column is the socket state (e.g. LISTEN); the address precedes it. + return listPortOccupation(ports, "lsof -iTCP -sTCP:LISTEN,TIME_WAIT -P -n", 10, 8, 1); } private static Pair getClusterNodesNum(final int index) { diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java new file mode 100644 index 000000000000..dbb9a9d7a74c --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java @@ -0,0 +1,47 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.it.env.cluster.EnvUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; + +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.net.InetAddress; +import java.net.ServerSocket; +import java.util.Collections; +import java.util.Map; + +import static org.junit.Assert.assertEquals; + +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class IoTDBTestEnvPortIT { + + @Test + public void testListeningPortIsReportedAsOccupied() throws Exception { + try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getByName("127.0.0.1"))) { + Map occupiedPorts = + EnvUtils.listPortOccupation(Collections.singletonList(server.getLocalPort())); + assertEquals( + Long.valueOf(ProcessHandle.current().pid()), occupiedPorts.get(server.getLocalPort())); + } + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java index be7cf49239a6..82fd7cc80e22 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java @@ -55,7 +55,8 @@ public class ClientManagerTest { - private final TEndPoint endPoint = new TEndPoint("localhost", 10730); + private final TEndPoint endPoint = + new TEndPoint("localhost", Integer.getInteger("iotdb.test.client-manager.port", 10730)); private static final int CONNECTION_TIMEOUT = 5_000; From 377bfc7acb5c2d8d6bb998ffc498113cdd12363b Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 14:20:40 +0800 Subject: [PATCH 06/15] Limit shared exchange budgets to eligible Collect fragments Retain independent plan-node quotas for dependent operators and internal pipelines, and fall back when a split side cannot hold one TsBlock. Make memory admission atomic without speculative charges and add regression coverage for MergeSort, queue forwarding, and contention. --- ...2026-10-14-exchange-fanin-budget-design.md | 47 ++- .../exchange/SharedTsBlockQueue.java | 3 +- .../exchange/sink/ShuffleSinkHandle.java | 4 +- .../execution/exchange/sink/SinkChannel.java | 3 +- .../exchange/source/SourceHandle.java | 3 +- .../execution/memory/MemoryPool.java | 99 +++-- .../DataNodeTableOperatorGenerator.java | 4 +- .../planner/LocalExecutionPlanContext.java | 64 ++-- .../plan/planner/LocalExecutionPlanner.java | 19 +- .../plan/planner/OperatorTreeGenerator.java | 7 +- .../SharedExchangeMemoryEligibility.java | 65 ++++ .../exchange/SharedTsBlockQueueTest.java | 30 ++ .../execution/memory/MemoryPoolTest.java | 88 +++++ .../LocalExecutionPlanContextBudgetTest.java | 356 ++++++++++++++++-- 14 files changed, 652 insertions(+), 140 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index e191622940dc..bc0b5c2557b2 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -26,21 +26,25 @@ ## 已确认范围 -实现 FI 输入输出共享账目、Collect/MappingCollect 任意就绪消费，以及由无序汇集引起的计划属性修正。 +实现仅针对无输入依赖 FI 的输入输出共享账目、Collect/MappingCollect 任意就绪消费，以及由无序汇集引起的计划属性修正。 聚合使用 HashAggregation：只清空实际输入已失去连续性的 preGroupedSymbols；有效分支内的流式聚合保留。 -MergeSort、Join、ASOF、CrossJoin 只作依赖验证，不修改其消费算法，不新增进度保障机制。 +MergeSort、Join、ASOF、CrossJoin 不修改消费算法。FI 内只要存在需要特定输入推进的算子，就整组使用原有按 plan node 限额；Collect 下方或上方出现这些算子也不例外。 TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍需匹配实际输入。 ## 额度与账目 -F 为配置的单 FI 交换预算。INPUT=F/2，OUTPUT=F-INPUT；同 FI 的所有 pipeline 继承相同固定值。 +共享采用 FI 级显式启用，默认保留原有配额。整个 FI 必须包含 CollectOperator 或 MappingCollectOperator，且所有多输入汇集都能消费任意就绪输入；顺序消费的 IdentitySink/CountMerge、MergeSort、Join、TopK 及未明确支持的多输入算子都禁止共享。检查在 Exchange 处停止，不跨到其他 FI。 + +在生成全部 pipeline 后再决定预算模式。内部 pipeline 队列具有转发依赖，因此存在多个 pipeline 或调度依赖时整组回退原策略。只含 Collect 的某个子树不会绕过 FI 级否决。 + +通过上述检查后，F 为配置的单 FI 交换预算，INPUT=F/2，OUTPUT=F-INPUT。两侧均须容纳配置的最大 TsBlock，否则回退原配额，避免仅有输出的 FI 因预算减半而无法打开 SinkChannel。未启用共享的 FI 恢复原有 exchangeSumNum 配额计算和按 plan node 的申请上限；不能只恢复限额数值却仍按同侧总量进行准入。 远端 SourceHandle 计入接收 FI 的 INPUT；SinkChannel 和 SharedTsBlockQueue 计入生产 FI 的 OUTPUT。 本地共享队列只计费一次，创建先后不改变所有者。 MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。 -FI 两侧使用 AtomicLong；planNode 字节明细通过 compute 增减。 -申请、回滚、释放均对应修改全局池、FI 单侧与 planNode 三层账目。 -额度判断使用本次增加后的单侧总数；稳定快照中的总数等于同侧明细之和。 +FI 两侧保留 AtomicLong 统计，plan node 保留字节明细。申请和释放在 FI 的记账锁内更新；全局池采用有界原子申请。申请前先检查适用的同侧总量或 plan node 限额，失败不产生临时计费，也不需要回滚，避免并发请求分别撞到不同限额后全部挂起。 + +释放后在 FI 记账锁外完成等待 Future；保留入队后的重试和每个 Future 的完成/取消同步。稳定快照中的 FI 同侧总量等于同侧明细之和。 保留现有最高优先级查询的零计费兜底与交换组件的收发、等待、释放机制。 ## Collect 生命周期 @@ -65,7 +69,8 @@ SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除 ## 验证 -覆盖共享额度并发、INPUT/OUTPUT 隔离、所有者、奇数 F、超大申请、回滚释放取消及稳定账目一致性； +覆盖纯 Collect/MappingCollect 显式启用、混合 FI 全组回退、Exchange 边界、内部 pipeline 否决、单块预算边界，以及共享/原配额 FI 共存； +覆盖共享额度并发、INPUT/OUTPUT 隔离、所有者、奇数 F、超大申请、失败准入无计费、释放取消及稳定账目一致性； 覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射； 覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。 运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。 @@ -117,3 +122,31 @@ Collect 后 TopK 不误用有序输入，以及有效分支保留流式聚合。 重新执行 `mvn -o install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests` 后运行前述两条 verify 命令；IoTDBCollectFanInIT 的 7 项测试在单节点和三 DataNode 集群中分别全部通过。 Spotless、Checkstyle、许可证与 `git diff --check` 检查通过。 + +## FI 共享范围修正验证（2026-10-15） + +共享改为完整 FI 的显式启用：仅支持 Collect/MappingCollect 的任意就绪汇集； +同一 FI 内的多输入依赖算子、内部 pipeline 队列或调度依赖会使整组回退。 +Exchange 边界之外的上游 FI 不参与本 FI 的判断。回退同时恢复 handle 限额和按 plan node 准入。 + +107 项定向单元测试通过，包括纯 Collect/Union、混合 FI、pipeline 依赖、单块额度边界、 +共享与独立配额共存、并发准入，以及真实 TableMergeSort 配合模拟输入和 MemoryPool 的回归。 +本地双队列转发回归也通过。Spotless、Checkstyle、许可证和 `git diff --check` 检查通过。 + +```sh +mvn -o test -pl iotdb-core/datanode -Dtest=LocalExecutionPlanContextBudgetTest,MemoryPoolTest,SharedTsBlockQueueTest,SourceHandleTest,SinkChannelTest,LocalSourceHandleTest,LocalSinkChannelTest,MPPDataExchangeManagerTest,PipelineBuilderTest,CollectInputPropertyTest +mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableSimpleIT -P with-integration-tests +mvn -o failsafe:integration-test@integration-test failsafe:verify@verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -P TableClusterIT -P with-integration-tests +``` + +单节点 verify 从当前源码构建依赖和测试节点；集群运行复用同一份节点模板。 +IoTDBCollectFanInIT 的 8 项测试在两种环境中均通过，零失败、零跳过。 + +提交前再次执行完整 reactor 的中英文编译，两者均通过： + +```sh +mvn -o test-compile -DskipTests -P with-zh-locale +mvn -o test-compile -DskipTests +``` + +本次生产代码没有新增未本地化的字符串，且没有修改中英文消息常量。 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java index cb8c681e4e9d..0d8423803618 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueue.java @@ -86,8 +86,7 @@ public class SharedTsBlockQueue { private LocalSinkChannel sinkChannel; private long maxBytesCanReserve = - MemorySide.OUTPUT.getMaxBytes( - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); private final boolean isHighestPriority; private volatile Throwable abortedCause = null; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java index 1de3ccb897df..e54062158a3f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/ShuffleSinkHandle.java @@ -22,7 +22,6 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager; -import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -77,8 +76,7 @@ public class ShuffleSinkHandle implements ISinkHandle { /** max bytes this ShuffleSinkHandle can reserve. */ private long maxBytesCanReserve = - MemorySide.OUTPUT.getMaxBytes( - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(ShuffleSinkHandle.class) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index 0dec1cbf07f2..4d592771ef23 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -132,8 +132,7 @@ private TsBlockInfo(TsBlock tsBlock, ByteBuffer serializedTsBlock, long tsBlockS /** max bytes this SinkChannel can reserve. */ private long maxBytesCanReserve = - MemorySide.OUTPUT.getMaxBytes( - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); private final boolean isHighestPriority; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index b1dc2781ed65..8ee7ad601ef6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -118,8 +118,7 @@ public class SourceHandle implements ISourceHandle { /** max bytes this SourceHandle can reserve. */ private long maxBytesCanReserve = - MemorySide.INPUT.getMaxBytes( - IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance()); + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); /** * this is set to true after calling isBlocked() at least once which indicates that this diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 07eced179d94..84f36b52e4ed 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -59,7 +59,8 @@ public static class MemoryReservationFuture extends AbstractFuture { /** * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the - * pool. This limit is shared by all plan nodes registered on the same side of that fragment. + * pool. For opted-in fragments the limit is shared by all plan nodes on the same side; + * otherwise it applies independently to the plan node. */ private final long maxBytesCanReserve; @@ -173,6 +174,8 @@ private static class FragmentMemoryReservation { private final Map planNodes = new ConcurrentHashMap<>(); private final AtomicLong inputUsedBytes = new AtomicLong(); private final AtomicLong outputUsedBytes = new AtomicLong(); + // Set by local planning before this FI starts reserving exchange memory. + private boolean sharedBudget; private AtomicLong getUsedBytes(MemorySide side) { return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; @@ -253,6 +256,19 @@ public void registerPlanNodeIdToQueryMemoryMap( } } + /** Opt in only after the complete FI has been checked for input and pipeline dependencies. */ + public void registerFragmentInstanceSharedBudget(String queryId, String fragmentInstanceId) { + synchronized (queryMemoryReservations) { + FragmentMemoryReservation fragment = + queryMemoryReservations + .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) + .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()); + synchronized (fragment) { + fragment.sharedBudget = true; + } + } + } + /** * If all fragmentInstanceIds related to one queryId have been registered, when the last fragment * instance is deregister, the queryId can be cleared. @@ -356,7 +372,6 @@ public MemoryReservationResult reserveWithPriority( if (tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve)) { return new MemoryReservationResult(immediateVoidFuture(), true, bytesToReserve); } else { - rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); if (isHighestPriority) { // SHOW QUERIES: treat as success with zero bytes reserved from pool when insufficient. return new MemoryReservationResult(immediateVoidFuture(), true, 0L); @@ -398,12 +413,7 @@ public boolean tryReserveForTest( maxBytesPerFragmentInstance, bytesToReserve); - if (tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve)) { - return true; - } else { - rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); - return false; - } + return tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); } /** @@ -447,20 +457,19 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l DataNodeQueryMessages .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2); } - fragment.planNodes.compute( - planNodeId, - (key, reservation) -> { - requireReservation(reservation, queryId, fragmentInstanceId, planNodeId); - if (reservation.bytesReserved < bytes) { - throw new IllegalArgumentException( - DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); - } - reservation.bytesReserved -= bytes; - fragment.getUsedBytes(reservation.side).addAndGet(-bytes); - return reservation; - }); - - memoryBlock.release(bytes); + synchronized (fragment) { + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + if (reservation.bytesReserved < bytes) { + throw new IllegalArgumentException( + DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); + } + reservation.bytesReserved -= bytes; + fragment.getUsedBytes(reservation.side).addAndGet(-bytes); + memoryBlock.release(bytes); + } + // Complete futures outside the accounting lock: listeners can reserve or free other inputs. if (memoryReservationFutures.isEmpty()) { return; @@ -492,7 +501,6 @@ private boolean tryCompleteReservation(MemoryReservationFuture future) { future.set(null); return true; } - rollbackReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve); return false; } } @@ -538,33 +546,24 @@ public boolean tryReserve( long maxBytesCanReserve) { FragmentMemoryReservation fragment = getFragmentReservation(queryId, fragmentInstanceId, planNodeId); - PlanNodeReservation reservation = - requireReservation( - fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); - long tryUsedBytes = memoryBlock.forceAllocateWithoutLimitation(bytesToReserve); - long sideUsedBytes = fragment.getUsedBytes(reservation.side).addAndGet(bytesToReserve); - fragment.planNodes.compute( - planNodeId, - (key, value) -> { - value.bytesReserved += bytesToReserve; - return value; - }); - return tryUsedBytes <= memoryBlock.getTotalMemorySizeInBytes() - && sideUsedBytes <= maxBytesCanReserve; - } - - private void rollbackReserve( - String queryId, String fragmentInstanceId, String planNodeId, long bytesToReserve) { - FragmentMemoryReservation fragment = - getFragmentReservation(queryId, fragmentInstanceId, planNodeId); - fragment.planNodes.compute( - planNodeId, - (key, reservation) -> { - reservation.bytesReserved -= bytesToReserve; - fragment.getUsedBytes(reservation.side).addAndGet(-bytesToReserve); - return reservation; - }); - memoryBlock.release(bytesToReserve); + synchronized (fragment) { + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + long usedBytes = + fragment.sharedBudget + ? fragment.getUsedBytes(reservation.side).get() + : reservation.bytesReserved; + // Do not publish speculative charges. Two competing requests must not reject one another + // against different counters and then both wait after their rollbacks restore capacity. + if (bytesToReserve > maxBytesCanReserve - usedBytes + || !memoryBlock.allocate(bytesToReserve)) { + return false; + } + fragment.getUsedBytes(reservation.side).addAndGet(bytesToReserve); + reservation.bytesReserved += bytesToReserve; + return true; + } } private FragmentMemoryReservation getFragmentReservation( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java index fd539be7467e..2fae9aa53708 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java @@ -343,8 +343,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte Operator child = node.getChildren().get(0).accept(this, context); List children = new ArrayList<>(1); children.add(child); - sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); - context.getDriverContext().setSink(sinkHandle); + context.setISink(sinkHandle); return new IdentitySinkOperator( operatorContext, children, downStreamChannelIndex, sinkHandle); } else { @@ -385,7 +384,6 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext if (!isSameNode) { context.addExchangeSumNum(1); } - sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index bfed28ccc4ca..82c58cad6559 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -32,9 +32,10 @@ import org.apache.iotdb.db.queryengine.execution.driver.DriverContext; import org.apache.iotdb.db.queryengine.execution.driver.SchemaDriverContext; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; -import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TemplatedInfo; @@ -43,6 +44,7 @@ import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.TimeValuePair; @@ -55,6 +57,7 @@ import java.time.ZoneId; import java.util.ArrayList; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -82,8 +85,8 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext private List pipelineDriverFactories; private List exchangeOperatorList = new ArrayList<>(); private int exchangeSumNum = 0; - private final long inputMaxBytes; - private final long outputMaxBytes; + private final long fragmentExchangeMaxBytes; + private Map legacySinkLimits = new IdentityHashMap<>(); private List cachedDataTypes; @@ -111,10 +114,8 @@ public LocalExecutionPlanContext( TypeProvider typeProvider, FragmentInstanceContext instanceContext, DataNodeQueryContext dataNodeQueryContext) { - long fragmentExchangeMaxBytes = + this.fragmentExchangeMaxBytes = IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); - this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); - this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); this.typeProvider = typeProvider; this.allSensorsMap = new ConcurrentHashMap<>(); this.nextOperatorId = new AtomicInteger(0); @@ -126,8 +127,8 @@ public LocalExecutionPlanContext( // For creating subContext, differ from parent context mainly in driver context public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) { - this.inputMaxBytes = parentContext.inputMaxBytes; - this.outputMaxBytes = parentContext.outputMaxBytes; + this.fragmentExchangeMaxBytes = parentContext.fragmentExchangeMaxBytes; + this.legacySinkLimits = parentContext.legacySinkLimits; this.nextOperatorId = parentContext.nextOperatorId; this.typeProvider = parentContext.typeProvider; this.allSensorsMap = parentContext.allSensorsMap; @@ -149,10 +150,8 @@ public LocalExecutionPlanContext( final TypeProvider typeProvider, final FragmentInstanceContext instanceContext, final ISchemaRegion schemaRegion) { - long fragmentExchangeMaxBytes = + this.fragmentExchangeMaxBytes = IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); - this.inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); - this.outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); this.allSensorsMap = new ConcurrentHashMap<>(); this.typeProvider = typeProvider; this.nextOperatorId = new AtomicInteger(0); @@ -258,12 +257,10 @@ public void setExchangeSumNum(int exchangeSumNum) { this.exchangeSumNum = exchangeSumNum; } - public long getInputMaxBytes() { - return inputMaxBytes; - } - - public long getOutputMaxBytes() { - return outputMaxBytes; + public long getMaxBytesOneHandleCanReserve() { + return exchangeSumNum == 0 + ? fragmentExchangeMaxBytes + : fragmentExchangeMaxBytes / exchangeSumNum; } public void addExchangeSumNum(int addValue) { @@ -274,16 +271,27 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { this.exchangeOperatorList.add(exchangeOperator); } - public void setExchangeMemoryLimits() { + public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { + long inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); + long outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); + long maxBlockBytes = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + boolean sharedBudget = + SharedExchangeMemoryEligibility.isEligible(plan) + // Internal pipeline queues and their downstream sinks must not compete for OUTPUT. + && pipelineDriverFactories.size() == 1 + && pipelineDriverFactories.get(0).getDependencyPipelineIndex() == -1 + && inputMaxBytes >= maxBlockBytes + && outputMaxBytes >= maxBlockBytes; + if (sharedBudget) { + FragmentInstanceId id = getFragmentInstanceId(); + memoryPool.registerFragmentInstanceSharedBudget(id.getQueryId().getId(), id.getFullId()); + } + legacySinkLimits.forEach( + (sink, legacyLimit) -> + sink.setMaxBytesCanReserve(sharedBudget ? outputMaxBytes : legacyLimit)); + long sourceLimit = sharedBudget ? inputMaxBytes : getMaxBytesOneHandleCanReserve(); exchangeOperatorList.forEach( - exchangeOperator -> { - exchangeOperator - .getSourceHandle() - .setMaxBytesCanReserve( - exchangeOperator.getSourceHandle() instanceof PipelineSourceHandle - ? outputMaxBytes - : inputMaxBytes); - }); + exchangeOperator -> exchangeOperator.getSourceHandle().setMaxBytesCanReserve(sourceLimit)); } public Set getAllSensors(IDeviceID deviceId, String sensorId) { @@ -316,7 +324,9 @@ public void setISink(ISink sink) { checkArgument( driverContext.getSink() == null, DataNodeQueryMessages.EXCEPTION_THERE_MUST_BE_AT_MOST_ONE_SINKNODE_A965AFE7); - sink.setMaxBytesCanReserve(outputMaxBytes); + if (sink instanceof ISinkHandle) { + legacySinkLimits.put(sink, getMaxBytesOneHandleCanReserve()); + } driverContext.setSink(sink); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index 1a85cb148c4f..8e020cf73645 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -35,6 +35,7 @@ import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.DeviceContext; import org.apache.iotdb.db.queryengine.execution.driver.DataDriverContext; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; @@ -132,8 +133,13 @@ public List plan( context.getTimePartitions().ifPresent(instanceContext::setTimePartitions); - // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setExchangeMemoryLimits(); + // Choose the accounting policy only after all FI operators and pipelines are known. + context.setExchangeMemoryLimits( + plan, + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool()); return context.getPipelineDriverFactories(); } @@ -165,8 +171,13 @@ public List plan( context.addPipelineDriverFactory(root, context.getDriverContext(), 0); - // set maxBytes one SourceHandle can reserve after visiting the whole tree - context.setExchangeMemoryLimits(); + // Choose the accounting policy only after all FI operators and pipelines are known. + context.setExchangeMemoryLimits( + plan, + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool()); return context.getPipelineDriverFactories(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java index d643ec03752e..a46877ac1d5f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java @@ -2577,7 +2577,6 @@ public Operator visitExchange(ExchangeNode node, LocalExecutionPlanContext conte if (!isSameNode) { context.addExchangeSumNum(1); } - sourceHandle.setMaxBytesCanReserve(context.getInputMaxBytes()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); @@ -2616,8 +2615,7 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte node.getPlanNodeId().getId(), context.getInstanceContext()); List children = dealWithConsumeChildrenOneByOneNode(node, context); - sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); - context.getDriverContext().setSink(sinkHandle); + context.setISink(sinkHandle); return new IdentitySinkOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); } @@ -2658,8 +2656,7 @@ public Operator visitShuffleSink(ShuffleSinkNode node, LocalExecutionPlanContext localInstanceId.toThrift(), node.getPlanNodeId().getId(), context.getInstanceContext()); - sinkHandle.setMaxBytesCanReserve(context.getOutputMaxBytes()); - context.getDriverContext().setSink(sinkHandle); + context.setISink(sinkHandle); return new ShuffleHelperOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java new file mode 100644 index 000000000000..f5a716a72b8f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java @@ -0,0 +1,65 @@ +/* + * 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.planner; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ExchangeNode; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.List; + +/** Select shared budgets only for FIs whose fan-in can consume any ready input. */ +final class SharedExchangeMemoryEligibility { + private SharedExchangeMemoryEligibility() {} + + static boolean isEligible(PlanNode root) { + boolean hasCollect = false; + Deque nodes = new ArrayDeque<>(); + nodes.push(root); + while (!nodes.isEmpty()) { + PlanNode node = nodes.pop(); + if (node instanceof ExchangeNode + || node + instanceof + org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode) { + // An upstream plan, if still attached, belongs to a different FI. + continue; + } + // Core relational nodes need not implement the serialization-only PlanNode.getType(). + boolean anyReady = + node instanceof CollectNode + || node instanceof UnionNode + || node + instanceof + org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.CollectNode; + hasCollect |= anyReady; + List children = node.getChildren(); + // Unknown multi-input operators are not implicitly safe. In particular, sequential sinks + // and joins can wait for an input while another input holds the entire shared budget. + if (!anyReady && children.size() > 1) { + return false; + } + nodes.addAll(children); + } + return hasCollect; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java index a92eaa5d0e28..8b97de5f3a25 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SharedTsBlockQueueTest.java @@ -254,4 +254,34 @@ public void run() { } } } + + @Test + public void testDependentQueuesKeepIndependentBudgetsByDefault() { + MemoryPool pool = new MemoryPool("pipeline-chain", new MemoryManager(1024), 512); + LocalMemoryManager manager = Mockito.mock(LocalMemoryManager.class); + Mockito.when(manager.getQueryPool()).thenReturn(pool); + ExecutorService executor = newDirectExecutorService(); + TFragmentInstanceId id = new TFragmentInstanceId("pipeline-chain", 0, "0"); + SharedTsBlockQueue upstream = new SharedTsBlockQueue(id, "a", manager, executor); + SharedTsBlockQueue downstream = new SharedTsBlockQueue(id, "b", manager, executor); + upstream.setMaxBytesCanReserve(256); + downstream.setMaxBytesCanReserve(256); + TsBlock block = Utils.createMockTsBlock(64); + try { + for (int i = 0; i < 4; i++) Assert.assertTrue(upstream.add(block).isDone()); + ListenableFuture fifth = upstream.add(block); + Assert.assertFalse(fifth.isDone()); + TsBlock forwarded = upstream.remove(); + Assert.assertTrue(fifth.isDone()); + Assert.assertTrue(downstream.add(forwarded).isDone()); + Assert.assertEquals(4, upstream.getNumOfBufferedTsBlocks()); + Assert.assertEquals(1, downstream.getNumOfBufferedTsBlocks()); + Assert.assertEquals(320, pool.getReservedBytes()); + } finally { + downstream.close(); + upstream.close(); + executor.shutdownNow(); + } + Assert.assertEquals(0, pool.getReservedBytes()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 0254fa1ac33c..c590e5a870da 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -57,6 +57,7 @@ public class MemoryPoolTest { public void before() { MemoryManager memoryManager = Mockito.spy(new MemoryManager(1024L)); pool = new MemoryPool("test", memoryManager, 512L); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } @@ -767,4 +768,91 @@ public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() t Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); } } + + @Test + public void testDefaultFragmentKeepsIndependentPlanNodeQuotas() { + MemoryPool legacy = new MemoryPool("legacy", new MemoryManager(1024), 512); + legacy.registerPlanNodeIdToQueryMemoryMap("q", "f", "a", MemorySide.INPUT); + legacy.registerPlanNodeIdToQueryMemoryMap("q", "f", "b", MemorySide.INPUT); + Assert.assertTrue( + legacy.reserveWithPriority("q", "f", "a", 256, 256, false).isReserveSuccess()); + Assert.assertTrue( + legacy.reserveWithPriority("q", "f", "b", 256, 256, false).isReserveSuccess()); + Assert.assertEquals(512, legacy.getFragmentMemoryReservedBytes("q", "f", MemorySide.INPUT)); + legacy.free("q", "f", "a", 256); + legacy.free("q", "f", "b", 256); + Assert.assertEquals(0, legacy.getReservedBytes()); + } + + @Test + public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "b", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false) + .isReserveSuccess()); + MemoryReservationResult blocked = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 64, 256, false); + Assert.assertFalse(blocked.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, "legacy", "a", 256, 256, false).isReserveSuccess()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, "legacy", "b", 256, 256, false).isReserveSuccess()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + Assert.assertTrue(blocked.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 64); + pool.free(QUERY_ID, "legacy", "a", 256); + pool.free(QUERY_ID, "legacy", "b", 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { + for (int i = 0; i < 3; i++) { + String fi = "other" + i; + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, fi, PLAN_NODE_ID, 256, 256, false).isReserveSuccess()); + } + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 100; i++) { + CountDownLatch start = new CountDownLatch(1); + Future first = + executor.submit( + () -> { + start.await(); + return pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 256, false); + }); + Future second = + executor.submit( + () -> { + start.await(); + return pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, "b", 256, 256, false); + }); + start.countDown(); + MemoryReservationResult a = first.get(10, TimeUnit.SECONDS); + MemoryReservationResult b = second.get(10, TimeUnit.SECONDS); + Assert.assertTrue(a.getFuture().isDone() ^ b.getFuture().isDone()); + String granted = a.getFuture().isDone() ? PLAN_NODE_ID : "b"; + String waiting = a.getFuture().isDone() ? "b" : PLAN_NODE_ID; + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, granted, 256); + Assert.assertTrue(a.getFuture().isDone()); + Assert.assertTrue(b.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, waiting, 256); + Assert.assertEquals(256, pool.getRemainingBytes()); + Assert.assertEquals( + 0, + pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + } + } finally { + executor.shutdownNow(); + } + for (int i = 0; i < 3; i++) pool.free(QUERY_ID, "other" + i, PLAN_NODE_ID, 256); + Assert.assertEquals(0, pool.getReservedBytes()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 68bff035fe2c..2d911fcd0189 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -19,60 +19,346 @@ package org.apache.iotdb.db.queryengine.plan.planner; +import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator; +import org.apache.iotdb.commons.memory.MemoryManager; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.common.QueryId; -import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; +import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; +import com.google.common.collect.ArrayListMultimap; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.junit.After; +import org.junit.Before; import org.junit.Test; -import static org.junit.Assert.assertEquals; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; public class LocalExecutionPlanContextBudgetTest { - @Test - public void testFixedBudgetsInheritedByPipelinesAndWiredBySide() { - long previous = + private long previousBudget; + private int previousBlockSize; + private MemoryPool pool; + private LocalExecutionPlanContext context; + + @Before + public void setUp() { + previousBudget = IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); - try { - IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(513); - FragmentInstanceContext instance = mock(FragmentInstanceContext.class); - when(instance.getId()) - .thenReturn(new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0")); - LocalExecutionPlanContext context = - new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); - context.setExchangeSumNum(500); - assertEquals(256, context.getInputMaxBytes()); - assertEquals(257, context.getOutputMaxBytes()); - IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); - LocalExecutionPlanContext child = context.createSubContext(); - assertEquals(256, child.getInputMaxBytes()); - assertEquals(257, child.getOutputMaxBytes()); - ISink sink = mock(ISink.class); - child.setISink(sink); - verify(sink).setMaxBytesCanReserve(257); - ISourceHandle remote = mock(ISourceHandle.class); - PipelineSourceHandle pipeline = mock(PipelineSourceHandle.class); - ExchangeOperator remoteOperator = mock(ExchangeOperator.class); - ExchangeOperator pipelineOperator = mock(ExchangeOperator.class); - when(remoteOperator.getSourceHandle()).thenReturn(remote); - when(pipelineOperator.getSourceHandle()).thenReturn(pipeline); - context.addExchangeOperator(remoteOperator); - child.addExchangeOperator(pipelineOperator); - context.setExchangeMemoryLimits(); - verify(remote).setMaxBytesCanReserve(256); - verify(pipeline).setMaxBytesCanReserve(257); - } finally { - IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previous); + previousBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(64); + pool = new MemoryPool("budget-test", new MemoryManager(4096), 1024); + context = newContext(); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previousBudget); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(previousBlockSize); + } + + private LocalExecutionPlanContext newContext() { + FragmentInstanceContext instance = mock(FragmentInstanceContext.class); + when(instance.getId()) + .thenReturn(new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0")); + LocalExecutionPlanContext result = + new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); + result.addPipelineDriverFactory(mock(Operator.class), result.getDriverContext(), 0); + result.setExchangeSumNum(4); + return result; + } + + private static PlanNode node(PlanNodeType type, PlanNode... children) { + Class nodeClass; + switch (type) { + case TABLE_EXCHANGE_NODE: + nodeClass = org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode.class; + break; + case EXCHANGE: + nodeClass = + org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ExchangeNode.class; + break; + case TABLE_COLLECT_NODE: + nodeClass = CollectNode.class; + break; + case TREE_COLLECT: + nodeClass = + org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.CollectNode.class; + break; + default: + nodeClass = PlanNode.class; + break; + } + PlanNode node = mock(nodeClass); + when(node.getType()).thenReturn(type); + when(node.getChildren()).thenReturn(Arrays.asList(children)); + return node; + } + + private static PlanNode collect() { + return new CollectNode( + new PlanNodeId("collect"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + Collections.emptyList()); + } + + private ISourceHandle addSource(LocalExecutionPlanContext target) { + ISourceHandle handle = mock(ISourceHandle.class); + ExchangeOperator operator = mock(ExchangeOperator.class); + when(operator.getSourceHandle()).thenReturn(handle); + target.addExchangeOperator(operator); + return handle; + } + + private void assertAdmissionIsShared(boolean shared) { + String fi = context.getFragmentInstanceId().getFullId(); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT); + assertTrue(pool.reserveWithPriority("budget", fi, "a", 256, 256, false).isReserveSuccess()); + MemoryPool.MemoryReservationResult second = + pool.reserveWithPriority("budget", fi, "b", 256, 256, false); + if (shared) { + assertFalse(second.getFuture().isDone()); + pool.free("budget", fi, "a", 256); + assertTrue(second.getFuture().isDone()); + } else { + assertTrue(second.isReserveSuccess()); + pool.free("budget", fi, "a", 256); + } + pool.free("budget", fi, "b", 256); + } + + @Test + public void testPureCollectUsesSharedSideBudgets() { + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(true); + } + + @Test + public void testMappingCollectUsesSharedBudget() { + PlanNode union = + new UnionNode( + new PlanNodeId("union"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + ArrayListMultimap.create(), + Collections.emptyList()); + context.setExchangeMemoryLimits(union, pool); + assertAdmissionIsShared(true); + } + + @Test + public void testTreeCollectUsesSharedBudget() { + context.setExchangeMemoryLimits( + node(PlanNodeType.TREE_COLLECT, node(PlanNodeType.EXCHANGE), node(PlanNodeType.EXCHANGE)), + pool); + assertAdmissionIsShared(true); + } + + @Test + public void testNestedDependencyDisablesSharingForEntireFragment() { + for (PlanNodeType dependency : + new PlanNodeType[] { + PlanNodeType.TABLE_MERGESORT_NODE, PlanNodeType.MERGE_SORT, PlanNodeType.AGG_MERGE_SORT, + PlanNodeType.TABLE_JOIN_NODE, PlanNodeType.TABLE_SEMI_JOIN_NODE, + PlanNodeType.TABLE_TOPK_NODE, + PlanNodeType.COUNT_MERGE, PlanNodeType.DEVICE_VIEW + }) { + PlanNode dependent = + node( + dependency, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE)); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.TABLE_COLLECT_NODE, dependent, collect()))); + assertFalse( + SharedExchangeMemoryEligibility.isEligible(node(dependency, collect(), collect()))); + } + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setExchangeSumNum(2); + context.setISink(sink); + context.setExchangeSumNum(4); + context.setExchangeMemoryLimits( + node( + PlanNodeType.TABLE_COLLECT_NODE, + collect(), + node( + PlanNodeType.TABLE_MERGESORT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE))), + pool); + verify(source).setMaxBytesCanReserve(256); + // The legacy sink quota is captured at its original point in operator generation. + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(false); + } + + @Test + public void testSequentialAndUnknownMultiInputOperatorsDisableSharing() { + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.IDENTITY_SINK, collect(), collect()))); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.SCHEMA_QUERY_MERGE, collect(), collect()))); + } + + @Test + public void testExchangeBoundaryDoesNotInspectUpstreamFragment() { + PlanNode remote = node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), collect()); + assertTrue( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE, remote), + node(PlanNodeType.TABLE_EXCHANGE_NODE)))); + } + + @Test + public void testNoCollectKeepsLegacyBudget() { + context.setExchangeMemoryLimits( + node( + PlanNodeType.TABLE_MERGESORT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE), + node(PlanNodeType.TABLE_EXCHANGE_NODE)), + pool); + assertAdmissionIsShared(false); + } + + @Test + public void testInternalPipelineDisablesSharingForEntireFragment() { + LocalExecutionPlanContext child = context.createSubContext(); + child.addPipelineDriverFactory(mock(Operator.class), child.getDriverContext(), 0); + ISourceHandle source = addSource(child); + PipelineSourceHandle pipelineSource = mock(PipelineSourceHandle.class); + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(pipelineSource); + context.addExchangeOperator(exchange); + context.setExchangeMemoryLimits(collect(), pool); + verify(source).setMaxBytesCanReserve(256); + verify(pipelineSource).setMaxBytesCanReserve(256); + assertAdmissionIsShared(false); + } + + @Test + public void testPipelineSchedulingDependencyDisablesSharing() { + context.getCurrentPipelineDriverFactory().setDependencyPipeline(0); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + + @Test + public void testSideSmallerThanOneBlockKeepsLegacyBudget() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + + @Test + public void testSinkOnlyFragmentCanStillReserveOneFullBlock() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(64); + context = newContext(); + context.setExchangeSumNum(1); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits( + node(PlanNodeType.IDENTITY_SINK, node(PlanNodeType.SERIES_SCAN)), pool); + verify(sink).setMaxBytesCanReserve(64); + } + + @Test + public void testMergeSortGetsBothHeadsDespitePrefetchOnOneInput() throws Exception { + context.setExchangeMemoryLimits( + node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + pool); + String fi = context.getFragmentInstanceId().getFullId(); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT); + for (int i = 0; i < 4; i++) { + assertTrue(pool.reserveWithPriority("budget", fi, "a", 64, 256, false).isReserveSuccess()); } + MemoryPool.MemoryReservationResult prefetch = + pool.reserveWithPriority("budget", fi, "a", 64, 256, false); + assertFalse(prefetch.getFuture().isDone()); + MemoryPool.MemoryReservationResult headB = + pool.reserveWithPriority("budget", fi, "b", 64, 256, false); + assertTrue(headB.isReserveSuccess()); + Operator a = mock(Operator.class); + Operator b = mock(Operator.class); + doReturn(Operator.NOT_BLOCKED).when(a).isBlocked(); + doReturn(headB.getFuture()).when(b).isBlocked(); + when(a.hasNextWithTimer()).thenReturn(true); + when(b.hasNextWithTimer()).thenReturn(true); + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)); + builder.getTimeColumnBuilder().writeLong(1); + builder.getColumnBuilder(0).writeInt(1); + builder.declarePosition(); + TsBlock first = builder.build(); + builder.reset(); + builder.getTimeColumnBuilder().writeLong(2); + builder.getColumnBuilder(0).writeInt(2); + builder.declarePosition(); + TsBlock second = builder.build(); + when(a.nextWithTimer()) + .thenAnswer( + call -> { + pool.free("budget", fi, "a", 64); + return first; + }); + when(b.nextWithTimer()) + .thenAnswer( + call -> { + pool.free("budget", fi, "b", 64); + return second; + }); + TableMergeSortOperator merge = + new TableMergeSortOperator( + mock(CommonOperatorContext.class), + new ArrayList<>(Arrays.asList(a, b)), + Collections.singletonList(TSDataType.INT32), + Comparator.comparingLong(key -> key.tsBlock.getTimeByIndex(key.rowIndex))); + assertTrue(merge.isBlocked().isDone()); + TsBlock result = merge.next(); + assertTrue(result != null && !result.isEmpty()); + assertTrue(prefetch.getFuture().isDone()); + pool.free("budget", fi, "a", 256); + merge.close(); } } From f53781cfa4a149c2e8eed062f1bd85c904535163 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 14:39:22 +0800 Subject: [PATCH 07/15] Preserve full exchange budgets for single-direction fragments Determine charged directions from remote sources and fragment sinks after planning. Split shared budgets only for bidirectional fragments and validate the maximum TsBlock size only for active directions. Add full-budget and local-input ownership regressions. --- ...2026-10-14-exchange-fanin-budget-design.md | 17 +++- .../execution/memory/MemorySide.java | 9 ++- .../planner/LocalExecutionPlanContext.java | 20 +++-- .../execution/memory/MemoryPoolTest.java | 29 ++++++- .../LocalExecutionPlanContextBudgetTest.java | 81 ++++++++++++++++++- 5 files changed, 144 insertions(+), 12 deletions(-) diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index bc0b5c2557b2..1c84700602d8 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -37,7 +37,7 @@ TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍 在生成全部 pipeline 后再决定预算模式。内部 pipeline 队列具有转发依赖，因此存在多个 pipeline 或调度依赖时整组回退原策略。只含 Collect 的某个子树不会绕过 FI 级否决。 -通过上述检查后，F 为配置的单 FI 交换预算，INPUT=F/2，OUTPUT=F-INPUT。两侧均须容纳配置的最大 TsBlock，否则回退原配额，避免仅有输出的 FI 因预算减半而无法打开 SinkChannel。未启用共享的 FI 恢复原有 exchangeSumNum 配额计算和按 plan node 的申请上限；不能只恢复限额数值却仍按同侧总量进行准入。 +通过上述检查后，F 为配置的单 FI 交换预算。按照实际由本 FI 计费的方向分配：只有 INPUT 时 INPUT=F；只有 OUTPUT 时 OUTPUT=F；同时存在两侧时 INPUT=F/2、OUTPUT=F-INPUT；未使用的方向预算为 0。方向在完整本地计划生成后确定，远端 SourceHandle 表示本 FI 的 INPUT，fragment sink 表示 OUTPUT；LocalSourceHandle 的队列由生产 FI 计费，不重复占用接收 FI 的 INPUT。仅实际启用的方向须容纳配置的最大 TsBlock，否则回退原配额。未启用共享的 FI 恢复原有 exchangeSumNum 配额计算和按 plan node 的申请上限；不能只恢复限额数值却仍按同侧总量进行准入。 远端 SourceHandle 计入接收 FI 的 INPUT；SinkChannel 和 SharedTsBlockQueue 计入生产 FI 的 OUTPUT。 本地共享队列只计费一次，创建先后不改变所有者。 @@ -150,3 +150,18 @@ mvn -o test-compile -DskipTests ``` 本次生产代码没有新增未本地化的字符串，且没有修改中英文消息常量。 + +## 单方向预算修正验证（2026-10-15） + +移除默认按两侧均分的单参数接口。`MemorySide.getMaxBytes(F, hasInput, hasOutput)` +显式接收本 FI 实际计费的方向：单方向使用完整 F，双方向才拆分，未使用方向为 0。 +最大 TsBlock 校验只作用于启用的方向；原有 Collect 资格和依赖回退规则保持不变。 + +新增输入单向、输出单向、B=F、本地输入和本地/远端混合输入用例； +同时验证奇数 F、未使用方向及每个单方向一次申请完整 F 的边界。 +113 项定向单测全部通过。首轮中单块边界测试的默认测试池上限为 512、申请为 513， +将该测试池配置为目标 F=513 后重跑 MemoryPoolTest，38 项全部通过，其余 75 项首轮通过。 + +完整 reactor 中英文 test-compile、Spotless、Checkstyle 和许可证检查通过。 +重新从当前代码打包节点后，IoTDBCollectFanInIT 在 TableSimpleIT 和三 DataNode 的 +TableClusterIT 中各 8 项全部通过；集群复用同一份节点模板，节点包哈希与本次构建一致。 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java index a59ef20d2684..d4c1908e2cc8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java @@ -24,7 +24,14 @@ public enum MemorySide { INPUT, OUTPUT; - public long getMaxBytes(long fragmentExchangeMaxBytes) { + /** Split only when both directions charge this FI; an unused direction has no budget. */ + public long getMaxBytes(long fragmentExchangeMaxBytes, boolean hasInput, boolean hasOutput) { + if (!(this == INPUT ? hasInput : hasOutput)) { + return 0; + } + if (!hasInput || !hasOutput) { + return fragmentExchangeMaxBytes; + } long inputMaxBytes = fragmentExchangeMaxBytes / 2; return this == INPUT ? inputMaxBytes : fragmentExchangeMaxBytes - inputMaxBytes; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index 82c58cad6559..4c0a516270c4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -33,6 +33,7 @@ import org.apache.iotdb.db.queryengine.execution.driver.SchemaDriverContext; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; @@ -272,16 +273,25 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { } public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { - long inputMaxBytes = MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes); - long outputMaxBytes = MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes); + // Only remote sources charge this FI's INPUT. Local inputs are owned by the producing FI; + // internal pipeline queues cannot enable sharing because of the pipeline check below. + boolean hasInput = + exchangeOperatorList.stream() + .anyMatch(exchange -> exchange.getSourceHandle() instanceof SourceHandle); + boolean hasOutput = !legacySinkLimits.isEmpty(); + long inputMaxBytes = + MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes, hasInput, hasOutput); + long outputMaxBytes = + MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes, hasInput, hasOutput); long maxBlockBytes = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); boolean sharedBudget = - SharedExchangeMemoryEligibility.isEligible(plan) + (hasInput || hasOutput) + && SharedExchangeMemoryEligibility.isEligible(plan) // Internal pipeline queues and their downstream sinks must not compete for OUTPUT. && pipelineDriverFactories.size() == 1 && pipelineDriverFactories.get(0).getDependencyPipelineIndex() == -1 - && inputMaxBytes >= maxBlockBytes - && outputMaxBytes >= maxBlockBytes; + && (!hasInput || inputMaxBytes >= maxBlockBytes) + && (!hasOutput || outputMaxBytes >= maxBlockBytes); if (sharedBudget) { FragmentInstanceId id = getFragmentInstanceId(); memoryPool.registerFragmentInstanceSharedBudget(id.getQueryId().getId(), id.getFullId()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index c590e5a870da..bf9b497ede88 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -577,8 +577,8 @@ public void testSideRegistrationCannotChange() { @Test public void testOddBudgetAndSingleRequestBoundary() { - Assert.assertEquals(256, MemorySide.INPUT.getMaxBytes(513)); - Assert.assertEquals(257, MemorySide.OUTPUT.getMaxBytes(513)); + Assert.assertEquals(256, MemorySide.INPUT.getMaxBytes(513, true, true)); + Assert.assertEquals(257, MemorySide.OUTPUT.getMaxBytes(513, true, true)); try { pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false); Assert.fail("request exceeds the input budget"); @@ -588,6 +588,31 @@ public void testOddBudgetAndSingleRequestBoundary() { } } + @Test + public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { + Assert.assertEquals(513, MemorySide.INPUT.getMaxBytes(513, true, false)); + Assert.assertEquals(0, MemorySide.OUTPUT.getMaxBytes(513, true, false)); + Assert.assertEquals(0, MemorySide.INPUT.getMaxBytes(513, false, true)); + Assert.assertEquals(513, MemorySide.OUTPUT.getMaxBytes(513, false, true)); + Assert.assertEquals(0, MemorySide.INPUT.getMaxBytes(513, false, false)); + Assert.assertEquals(0, MemorySide.OUTPUT.getMaxBytes(513, false, false)); + MemoryPool singleDirectionPool = + new MemoryPool("single-direction", new MemoryManager(1024), 513); + singleDirectionPool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + for (MemorySide side : MemorySide.values()) { + String node = side.name(); + singleDirectionPool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, node, side); + long limit = side.getMaxBytes(513, side == MemorySide.INPUT, side == MemorySide.OUTPUT); + Assert.assertTrue( + singleDirectionPool + .reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, node, 513, limit, false) + .isReserveSuccess()); + singleDirectionPool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, node, 513); + Assert.assertEquals(0, singleDirectionPool.getReservedBytes()); + } + } + @Test public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { ExecutorService executor = Executors.newFixedThreadPool(16); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 2d911fcd0189..2f1035245bfc 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -34,7 +34,9 @@ import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.PipelineSourceHandle; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; @@ -133,7 +135,7 @@ private static PlanNode collect() { } private ISourceHandle addSource(LocalExecutionPlanContext target) { - ISourceHandle handle = mock(ISourceHandle.class); + ISourceHandle handle = mock(SourceHandle.class); ExchangeOperator operator = mock(ExchangeOperator.class); when(operator.getSourceHandle()).thenReturn(handle); target.addExchangeOperator(operator); @@ -141,9 +143,13 @@ private ISourceHandle addSource(LocalExecutionPlanContext target) { } private void assertAdmissionIsShared(boolean shared) { + assertAdmissionIsShared(shared, MemorySide.INPUT); + } + + private void assertAdmissionIsShared(boolean shared, MemorySide side) { String fi = context.getFragmentInstanceId().getFullId(); - pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT); - pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", side); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", side); assertTrue(pool.reserveWithPriority("budget", fi, "a", 256, 256, false).isReserveSuccess()); MemoryPool.MemoryReservationResult second = pool.reserveWithPriority("budget", fi, "b", 256, 256, false); @@ -169,8 +175,74 @@ public void testPureCollectUsesSharedSideBudgets() { assertAdmissionIsShared(true); } + @Test + public void testInputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeSumNum(1); + ISourceHandle source = addSource(context); + context.setExchangeMemoryLimits( + node(PlanNodeType.TABLE_COLLECT_NODE, node(PlanNodeType.TABLE_EXCHANGE_NODE)), pool); + verify(source).setMaxBytesCanReserve(1024); + assertAdmissionIsShared(true); + } + + @Test + public void testOutputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + context.setExchangeSumNum(1); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits( + node( + PlanNodeType.IDENTITY_SINK, + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.SERIES_SCAN), + node(PlanNodeType.SERIES_SCAN))), + pool); + verify(sink).setMaxBytesCanReserve(1024); + assertAdmissionIsShared(true, MemorySide.OUTPUT); + } + + private void addLocalInput() { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(mock(LocalSourceHandle.class)); + context.addExchangeOperator(exchange); + } + + @Test + public void testLocalInputsDoNotSplitOutputBudget() { + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); + addLocalInput(); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(sink).setMaxBytesCanReserve(1024); + assertAdmissionIsShared(true, MemorySide.OUTPUT); + } + + @Test + public void testMixedLocalAndRemoteInputsStillSplitBidirectionalBudget() { + addLocalInput(); + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(512); + assertAdmissionIsShared(true); + } + + @Test + public void testLocalInputsWithoutSinkDoNotRegisterUnusedSharedBudget() { + addLocalInput(); + context.setExchangeMemoryLimits(collect(), pool); + assertAdmissionIsShared(false); + } + @Test public void testMappingCollectUsesSharedBudget() { + addSource(context); PlanNode union = new UnionNode( new PlanNodeId("union"), @@ -184,6 +256,7 @@ public void testMappingCollectUsesSharedBudget() { @Test public void testTreeCollectUsesSharedBudget() { + addSource(context); context.setExchangeMemoryLimits( node(PlanNodeType.TREE_COLLECT, node(PlanNodeType.EXCHANGE), node(PlanNodeType.EXCHANGE)), pool); @@ -286,6 +359,8 @@ public void testPipelineSchedulingDependencyDisablesSharing() { @Test public void testSideSmallerThanOneBlockKeepsLegacyBudget() { + addSource(context); + context.setISink(mock(ISinkHandle.class)); TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(1024); context.setExchangeMemoryLimits(collect(), pool); assertAdmissionIsShared(false); From 1641461b0a74262d8876e220cc00f7897ac76f27 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 15:16:31 +0800 Subject: [PATCH 08/15] Initialize shared exchange limits in fragment memory records --- ...2026-10-14-exchange-fanin-budget-design.md | 20 ++++- .../execution/memory/MemoryPool.java | 31 +++++++- .../execution/memory/MemorySide.java | 14 +--- .../planner/LocalExecutionPlanContext.java | 24 +++--- .../execution/memory/MemoryPoolTest.java | 76 ++++++++++++++----- .../LocalExecutionPlanContextBudgetTest.java | 40 +++++++--- 6 files changed, 148 insertions(+), 57 deletions(-) diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index 1c84700602d8..87dcfa5fb8c9 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -42,6 +42,7 @@ TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍 本地共享队列只计费一次，创建先后不改变所有者。 MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。 +本地计划生成全部算子和 pipeline 后，一次性检查共享资格并计算两个方向的预算，写入 FI 记录的 `inputMaxBytes`、`outputMaxBytes` 字段。`MemorySide` 仅表示方向，不计算预算，也不接受 `hasInput`/`hasOutput` 参数。共享模式的申请、单块上限检查和等待唤醒统一读取 FI 保存的预算；handle 传入的限额仅用于原有独立配额模式。 FI 两侧保留 AtomicLong 统计，plan node 保留字节明细。申请和释放在 FI 的记账锁内更新；全局池采用有界原子申请。申请前先检查适用的同侧总量或 plan node 限额，失败不产生临时计费，也不需要回滚，避免并发请求分别撞到不同限额后全部挂起。 释放后在 FI 记账锁外完成等待 Future；保留入队后的重试和每个 Future 的完成/取消同步。稳定快照中的 FI 同侧总量等于同侧明细之和。 @@ -153,8 +154,9 @@ mvn -o test-compile -DskipTests ## 单方向预算修正验证（2026-10-15） -移除默认按两侧均分的单参数接口。`MemorySide.getMaxBytes(F, hasInput, hasOutput)` -显式接收本 FI 实际计费的方向：单方向使用完整 F，双方向才拆分，未使用方向为 0。 +该轮修正取消默认按两侧均分，改为根据本 FI 实际计费方向分配： +单方向使用完整 F，双方向才拆分，未使用方向为 0。本次已将计算统一收回本地计划初始化， +不再保留该轮引入的 `MemorySide.getMaxBytes(F, hasInput, hasOutput)` 接口。 最大 TsBlock 校验只作用于启用的方向；原有 Collect 资格和依赖回退规则保持不变。 新增输入单向、输出单向、B=F、本地输入和本地/远端混合输入用例； @@ -165,3 +167,17 @@ mvn -o test-compile -DskipTests 完整 reactor 中英文 test-compile、Spotless、Checkstyle 和许可证检查通过。 重新从当前代码打包节点后，IoTDBCollectFanInIT 在 TableSimpleIT 和三 DataNode 的 TableClusterIT 中各 8 项全部通过；集群复用同一份节点模板，节点包哈希与本次构建一致。 + +## FI 预算字段初始化验证（2026-10-15） + +`FragmentMemoryReservation` 保存 `inputMaxBytes`、`outputMaxBytes`； +完整本地计划生成后按输入 exchange、fragment sink 和依赖检查结果一次性初始化。 +共享申请、单块上限检查与等待重试均读取 FI 字段，独立配额模式仍读取 handle 限额。 +`MemorySide` 只保留方向枚举，移除布尔参数预算接口。 + +115 项定向单测通过，覆盖 handle 参数不能增减共享上限、超额请求不进入等待队列、 +未使用方向预算为零、奇数预算拆分和原有独立配额回退。 +完整 reactor 中英文编译、Spotless、Checkstyle 和许可证检查通过；生产代码没有新增字符串字面量。 + +重新打包当前源码后，`IoTDBCollectFanInIT` 在 TableSimpleIT 和三 DataNode 的 TableClusterIT +中各 8 项全部通过，无失败或跳过。集群复用同一份节点模板，其服务端 jar 哈希与本次构建一致。 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 84f36b52e4ed..98108b324758 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -59,8 +59,8 @@ public static class MemoryReservationFuture extends AbstractFuture { /** * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the - * pool. For opted-in fragments the limit is shared by all plan nodes on the same side; - * otherwise it applies independently to the plan node. + * pool. This is the per-plan-node limit for legacy fragments. Shared fragments always use their + * initialized FI-side limit, including when retrying a pending reservation. */ private final long maxBytesCanReserve; @@ -176,6 +176,14 @@ private static class FragmentMemoryReservation { private final AtomicLong outputUsedBytes = new AtomicLong(); // Set by local planning before this FI starts reserving exchange memory. private boolean sharedBudget; + private long inputMaxBytes; + private long outputMaxBytes; + + private long getMaxBytesCanReserve(PlanNodeReservation reservation, long legacyLimit) { + return sharedBudget + ? (reservation.side == MemorySide.INPUT ? inputMaxBytes : outputMaxBytes) + : legacyLimit; + } private AtomicLong getUsedBytes(MemorySide side) { return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; @@ -256,14 +264,20 @@ public void registerPlanNodeIdToQueryMemoryMap( } } - /** Opt in only after the complete FI has been checked for input and pipeline dependencies. */ - public void registerFragmentInstanceSharedBudget(String queryId, String fragmentInstanceId) { + /** + * Initialize once after the complete FI has been checked for input and pipeline dependencies, + * before any driver starts reserving exchange memory. An unused direction has a zero budget. + */ + public void registerFragmentInstanceSharedBudget( + String queryId, String fragmentInstanceId, long inputMaxBytes, long outputMaxBytes) { synchronized (queryMemoryReservations) { FragmentMemoryReservation fragment = queryMemoryReservations .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()); synchronized (fragment) { + fragment.inputMaxBytes = inputMaxBytes; + fragment.outputMaxBytes = outputMaxBytes; fragment.sharedBudget = true; } } @@ -354,6 +368,14 @@ public MemoryReservationResult reserveWithPriority( .EXCEPTION_BYTESTORESERVE_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTESPERFI_RIGHT_BRACKET_DOT_M_0753BB69, maxBytesPerFragmentInstance, bytesToReserve); + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + synchronized (fragment) { + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + maxBytesCanReserve = fragment.getMaxBytesCanReserve(reservation, maxBytesCanReserve); + } if (bytesToReserve > maxBytesCanReserve) { LOGGER.warn( DataNodeQueryMessages @@ -550,6 +572,7 @@ public boolean tryReserve( PlanNodeReservation reservation = requireReservation( fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + maxBytesCanReserve = fragment.getMaxBytesCanReserve(reservation, maxBytesCanReserve); long usedBytes = fragment.sharedBudget ? fragment.getUsedBytes(reservation.side).get() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java index d4c1908e2cc8..714996053855 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemorySide.java @@ -22,17 +22,5 @@ /** The exchange budget charged by a registered plan node. */ public enum MemorySide { INPUT, - OUTPUT; - - /** Split only when both directions charge this FI; an unused direction has no budget. */ - public long getMaxBytes(long fragmentExchangeMaxBytes, boolean hasInput, boolean hasOutput) { - if (!(this == INPUT ? hasInput : hasOutput)) { - return 0; - } - if (!hasInput || !hasOutput) { - return fragmentExchangeMaxBytes; - } - long inputMaxBytes = fragmentExchangeMaxBytes / 2; - return this == INPUT ? inputMaxBytes : fragmentExchangeMaxBytes - inputMaxBytes; - } + OUTPUT } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index 4c0a516270c4..e78240ec0bce 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -37,7 +37,6 @@ import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; -import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TemplatedInfo; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; @@ -275,26 +274,29 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { // Only remote sources charge this FI's INPUT. Local inputs are owned by the producing FI; // internal pipeline queues cannot enable sharing because of the pipeline check below. - boolean hasInput = + long inputExchangeCount = exchangeOperatorList.stream() - .anyMatch(exchange -> exchange.getSourceHandle() instanceof SourceHandle); - boolean hasOutput = !legacySinkLimits.isEmpty(); + .filter(exchange -> exchange.getSourceHandle() instanceof SourceHandle) + .count(); + int outputSinkCount = legacySinkLimits.size(); long inputMaxBytes = - MemorySide.INPUT.getMaxBytes(fragmentExchangeMaxBytes, hasInput, hasOutput); - long outputMaxBytes = - MemorySide.OUTPUT.getMaxBytes(fragmentExchangeMaxBytes, hasInput, hasOutput); + inputExchangeCount == 0 + ? 0 + : (outputSinkCount == 0 ? fragmentExchangeMaxBytes : fragmentExchangeMaxBytes / 2); + long outputMaxBytes = outputSinkCount == 0 ? 0 : fragmentExchangeMaxBytes - inputMaxBytes; long maxBlockBytes = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); boolean sharedBudget = - (hasInput || hasOutput) + (inputMaxBytes > 0 || outputMaxBytes > 0) && SharedExchangeMemoryEligibility.isEligible(plan) // Internal pipeline queues and their downstream sinks must not compete for OUTPUT. && pipelineDriverFactories.size() == 1 && pipelineDriverFactories.get(0).getDependencyPipelineIndex() == -1 - && (!hasInput || inputMaxBytes >= maxBlockBytes) - && (!hasOutput || outputMaxBytes >= maxBlockBytes); + && (inputExchangeCount == 0 || inputMaxBytes >= maxBlockBytes) + && (outputSinkCount == 0 || outputMaxBytes >= maxBlockBytes); if (sharedBudget) { FragmentInstanceId id = getFragmentInstanceId(); - memoryPool.registerFragmentInstanceSharedBudget(id.getQueryId().getId(), id.getFullId()); + memoryPool.registerFragmentInstanceSharedBudget( + id.getQueryId().getId(), id.getFullId(), inputMaxBytes, outputMaxBytes); } legacySinkLimits.forEach( (sink, legacyLimit) -> diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index bf9b497ede88..136e317e665b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -57,7 +57,6 @@ public class MemoryPoolTest { public void before() { MemoryManager memoryManager = Mockito.spy(new MemoryManager(1024L)); pool = new MemoryPool("test", memoryManager, 512L); - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } @@ -530,6 +529,7 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { @Test public void testSharedInputBudgetAndIndependentOutput() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -576,45 +576,80 @@ public void testSideRegistrationCannotChange() { } @Test - public void testOddBudgetAndSingleRequestBoundary() { - Assert.assertEquals(256, MemorySide.INPUT.getMaxBytes(513, true, true)); - Assert.assertEquals(257, MemorySide.OUTPUT.getMaxBytes(513, true, true)); + public void testSharedRequestCannotExceedInitializedInputBudget() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); try { - pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false); - Assert.fail("request exceeds the input budget"); + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, Long.MAX_VALUE, false); + Assert.fail("request exceeds the initialized input budget"); } catch (IllegalArgumentException expected) { Assert.assertEquals(0, pool.getReservedBytes()); Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); + Assert.assertEquals(0, pool.getMemoryReservationSize()); } } + @Test + public void testHandleLimitsCannotChangeInitializedSharedBudgets() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 1, false) + .isReserveSuccess()); + MemoryReservationResult pending = + pool.reserveWithPriority( + QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, Long.MAX_VALUE, false); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 1, false) + .isReserveSuccess()); + // Waking the input request must still respect the FI's saved limit. + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256); + Assert.assertFalse(pending.getFuture().isDone()); + Assert.assertFalse( + pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, Long.MAX_VALUE)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200); + pending.getFuture().get(10, TimeUnit.SECONDS); + Assert.assertEquals( + 100, pool.getFragmentMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100); + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } + @Test public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { - Assert.assertEquals(513, MemorySide.INPUT.getMaxBytes(513, true, false)); - Assert.assertEquals(0, MemorySide.OUTPUT.getMaxBytes(513, true, false)); - Assert.assertEquals(0, MemorySide.INPUT.getMaxBytes(513, false, true)); - Assert.assertEquals(513, MemorySide.OUTPUT.getMaxBytes(513, false, true)); - Assert.assertEquals(0, MemorySide.INPUT.getMaxBytes(513, false, false)); - Assert.assertEquals(0, MemorySide.OUTPUT.getMaxBytes(513, false, false)); MemoryPool singleDirectionPool = new MemoryPool("single-direction", new MemoryManager(1024), 513); - singleDirectionPool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); for (MemorySide side : MemorySide.values()) { - String node = side.name(); + String fi = side.name(); + singleDirectionPool.registerFragmentInstanceSharedBudget( + QUERY_ID, fi, side == MemorySide.INPUT ? 513 : 0, side == MemorySide.OUTPUT ? 513 : 0); + singleDirectionPool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, "active", side); singleDirectionPool.registerPlanNodeIdToQueryMemoryMap( - QUERY_ID, FRAGMENT_INSTANCE_ID, node, side); - long limit = side.getMaxBytes(513, side == MemorySide.INPUT, side == MemorySide.OUTPUT); + QUERY_ID, fi, "unused", side == MemorySide.INPUT ? MemorySide.OUTPUT : MemorySide.INPUT); Assert.assertTrue( singleDirectionPool - .reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, node, 513, limit, false) + .reserveWithPriority(QUERY_ID, fi, "active", 513, 1, false) .isReserveSuccess()); - singleDirectionPool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, node, 513); + try { + singleDirectionPool.reserveWithPriority(QUERY_ID, fi, "unused", 1, Long.MAX_VALUE, false); + Assert.fail("an unused direction must not acquire a budget from a handle"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, singleDirectionPool.getMemoryReservationSize()); + Assert.assertEquals(513, singleDirectionPool.getReservedBytes()); + } + singleDirectionPool.free(QUERY_ID, fi, "active", 513); Assert.assertEquals(0, singleDirectionPool.getReservedBytes()); } } @Test public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 128, 128); ExecutorService executor = Executors.newFixedThreadPool(16); CountDownLatch start = new CountDownLatch(1); CountDownLatch attempted = new CountDownLatch(16); @@ -662,6 +697,7 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { @Test public void testReleaseBeforePendingReservationIsPublished() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); pool = Mockito.spy(pool); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); @@ -714,6 +750,7 @@ public void testReleaseBeforePendingReservationIsPublished() throws Exception { @Test public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -751,6 +788,7 @@ public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Except @Test public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 128, 128); int count = 16; ExecutorService executor = Executors.newFixedThreadPool(count); CountDownLatch attempted = new CountDownLatch(count); @@ -811,6 +849,7 @@ public void testDefaultFragmentKeepsIndependentPlanNodeQuotas() { @Test public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "a", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "b", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); @@ -834,6 +873,7 @@ public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { @Test public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); for (int i = 0; i < 3; i++) { String fi = "other" + i; pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 2f1035245bfc..7d2e23e9d554 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -147,21 +147,26 @@ private void assertAdmissionIsShared(boolean shared) { } private void assertAdmissionIsShared(boolean shared, MemorySide side) { + assertAdmissionIsShared(shared, side, shared ? 1024 : 256); + } + + private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit) { String fi = context.getFragmentInstanceId().getFullId(); - pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", side); - pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", side); - assertTrue(pool.reserveWithPriority("budget", fi, "a", 256, 256, false).isReserveSuccess()); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a" + side, side); + pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b" + side, side); + assertTrue( + pool.reserveWithPriority("budget", fi, "a" + side, limit, limit, false).isReserveSuccess()); MemoryPool.MemoryReservationResult second = - pool.reserveWithPriority("budget", fi, "b", 256, 256, false); + pool.reserveWithPriority("budget", fi, "b" + side, limit, limit, false); if (shared) { assertFalse(second.getFuture().isDone()); - pool.free("budget", fi, "a", 256); + pool.free("budget", fi, "a" + side, limit); assertTrue(second.getFuture().isDone()); } else { assertTrue(second.isReserveSuccess()); - pool.free("budget", fi, "a", 256); + pool.free("budget", fi, "a" + side, limit); } - pool.free("budget", fi, "b", 256); + pool.free("budget", fi, "b" + side, limit); } @Test @@ -172,7 +177,23 @@ public void testPureCollectUsesSharedSideBudgets() { context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); verify(source).setMaxBytesCanReserve(512); verify(sink).setMaxBytesCanReserve(512); - assertAdmissionIsShared(true); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 512); + } + + @Test + public void testOddBudgetIsInitializedAfterBothDirectionsAreKnown() { + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1025); + pool = new MemoryPool("odd-budget", new MemoryManager(4096), 1025); + context = newContext(); + ISourceHandle source = addSource(context); + ISinkHandle sink = mock(ISinkHandle.class); + context.setISink(sink); + context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); + verify(source).setMaxBytesCanReserve(512); + verify(sink).setMaxBytesCanReserve(513); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 513); } @Test @@ -230,7 +251,8 @@ public void testMixedLocalAndRemoteInputsStillSplitBidirectionalBudget() { context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); verify(source).setMaxBytesCanReserve(512); verify(sink).setMaxBytesCanReserve(512); - assertAdmissionIsShared(true); + assertAdmissionIsShared(true, MemorySide.INPUT, 512); + assertAdmissionIsShared(true, MemorySide.OUTPUT, 512); } @Test From 21a42d419fbf40300c57fe1e5a6a5abaea1341ec Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 15:42:07 +0800 Subject: [PATCH 09/15] Move exchange budget configuration into fragment instance context --- ...2026-10-14-exchange-fanin-budget-design.md | 17 ++++++- .../fragment/FragmentInstanceContext.java | 18 ++++++++ .../execution/memory/MemoryPool.java | 33 +++++++------- .../planner/LocalExecutionPlanContext.java | 7 +-- .../execution/memory/MemoryPoolTest.java | 44 ++++++++++++++----- .../LocalExecutionPlanContextBudgetTest.java | 23 ++++++++-- 6 files changed, 108 insertions(+), 34 deletions(-) diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md index 87dcfa5fb8c9..3eb5cbc90625 100644 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md @@ -42,7 +42,7 @@ TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍 本地共享队列只计费一次，创建先后不改变所有者。 MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。 -本地计划生成全部算子和 pipeline 后，一次性检查共享资格并计算两个方向的预算，写入 FI 记录的 `inputMaxBytes`、`outputMaxBytes` 字段。`MemorySide` 仅表示方向，不计算预算，也不接受 `hasInput`/`hasOutput` 参数。共享模式的申请、单块上限检查和等待唤醒统一读取 FI 保存的预算；handle 传入的限额仅用于原有独立配额模式。 +本地计划生成全部算子和 pipeline 后，一次性检查共享资格并计算两个方向的预算，写入 `FragmentInstanceContext` 的 `inputMaxBytes`、`outputMaxBytes` 字段。`MemorySide` 仅表示方向，不计算预算，也不接受 `hasInput`/`hasOutput` 参数。MemoryPool 的 FI 账目仅保留该上下文的引用，不复制预算字段。共享模式的申请、单块上限检查和等待唤醒统一读取 FI 保存的预算；handle 传入的限额仅用于原有独立配额模式。回退独立配额时上下文的两个共享预算字段为零。 FI 两侧保留 AtomicLong 统计，plan node 保留字节明细。申请和释放在 FI 的记账锁内更新；全局池采用有界原子申请。申请前先检查适用的同侧总量或 plan node 限额，失败不产生临时计费，也不需要回滚，避免并发请求分别撞到不同限额后全部挂起。 释放后在 FI 记账锁外完成等待 Future；保留入队后的重试和每个 Future 的完成/取消同步。稳定快照中的 FI 同侧总量等于同侧明细之和。 @@ -181,3 +181,18 @@ TableClusterIT 中各 8 项全部通过；集群复用同一份节点模板， 重新打包当前源码后，`IoTDBCollectFanInIT` 在 TableSimpleIT 和三 DataNode 的 TableClusterIT 中各 8 项全部通过，无失败或跳过。集群复用同一份节点模板，其服务端 jar 哈希与本次构建一致。 + +## FI 预算配置归属修正验证（2026-10-15） + +预算字段移至 `FragmentInstanceContext`，由完整本地计划生成后的初始化流程赋值。 +MemoryPool 的 FI 账目保存上下文引用，通过 getter 读取限额，不再复制预算或保存共享布尔字段。 +该引用随 FI 记账记录注销；预算仍在 driver 启动前初始化。 +独立配额回退时上下文的两个共享预算字段为零，不注册共享预算。 + +123 项定向单测通过，包括真实 FI 上下文的初始化字段断言、申请上限和等待唤醒回归， +以及已有的 FI 上下文生命周期测试。生产代码没有新增字符串字面量。 + +本次单测在前述 `-Dtest` 列表中加入 `FragmentInstanceContextTest`。完整 reactor 的中英文 +`test-compile -DskipTests`、Spotless、Checkstyle、许可证与 `git diff --check` 检查均通过。 +重新打包节点后，`IoTDBCollectFanInIT` 在 TableSimpleIT 和三 DataNode 的 TableClusterIT 中各 +8 项全部通过，无失败或跳过；测试模板与当前服务端 jar 的哈希一致。 diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 5b8ee6ae1c1a..16fecf6bc498 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -120,6 +120,11 @@ public class FragmentInstanceContext extends QueryContext { private final MemoryReservationManager memoryReservationManager; + // Initialized after the complete local plan is generated, before drivers start. Both are zero + // when this FI uses independent per-plan-node exchange quotas. + private long inputMaxBytes; + private long outputMaxBytes; + protected IDataRegionForQuery dataRegion; private Filter globalTimeFilter; private List globalTimeFilterTimeRanges; @@ -551,6 +556,19 @@ public FragmentInstanceId getId() { return id; } + public void initializeExchangeMemoryBudget(long inputMaxBytes, long outputMaxBytes) { + this.inputMaxBytes = inputMaxBytes; + this.outputMaxBytes = outputMaxBytes; + } + + public long getInputMaxBytes() { + return inputMaxBytes; + } + + public long getOutputMaxBytes() { + return outputMaxBytes; + } + public void failed(Throwable cause) { stateMachine.failed(cause); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 98108b324758..8c3ebcbbc9fb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.exception.runtime.MemoryLeakException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import com.google.common.util.concurrent.AbstractFuture; import com.google.common.util.concurrent.ListenableFuture; @@ -174,14 +175,18 @@ private static class FragmentMemoryReservation { private final Map planNodes = new ConcurrentHashMap<>(); private final AtomicLong inputUsedBytes = new AtomicLong(); private final AtomicLong outputUsedBytes = new AtomicLong(); - // Set by local planning before this FI starts reserving exchange memory. - private boolean sharedBudget; - private long inputMaxBytes; - private long outputMaxBytes; + // Registered after local planning; the FI context owns the initialized shared limits. + private FragmentInstanceContext instanceContext; + + private boolean isSharedBudget() { + return instanceContext != null; + } private long getMaxBytesCanReserve(PlanNodeReservation reservation, long legacyLimit) { - return sharedBudget - ? (reservation.side == MemorySide.INPUT ? inputMaxBytes : outputMaxBytes) + return isSharedBudget() + ? (reservation.side == MemorySide.INPUT + ? instanceContext.getInputMaxBytes() + : instanceContext.getOutputMaxBytes()) : legacyLimit; } @@ -264,21 +269,17 @@ public void registerPlanNodeIdToQueryMemoryMap( } } - /** - * Initialize once after the complete FI has been checked for input and pipeline dependencies, - * before any driver starts reserving exchange memory. An unused direction has a zero budget. - */ - public void registerFragmentInstanceSharedBudget( - String queryId, String fragmentInstanceId, long inputMaxBytes, long outputMaxBytes) { + /** Register the FI-owned shared limits after local planning and before drivers start. */ + public void registerFragmentInstanceSharedBudget(FragmentInstanceContext instanceContext) { + String queryId = instanceContext.getId().getQueryId().getId(); + String fragmentInstanceId = instanceContext.getId().getFullId(); synchronized (queryMemoryReservations) { FragmentMemoryReservation fragment = queryMemoryReservations .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()); synchronized (fragment) { - fragment.inputMaxBytes = inputMaxBytes; - fragment.outputMaxBytes = outputMaxBytes; - fragment.sharedBudget = true; + fragment.instanceContext = instanceContext; } } } @@ -574,7 +575,7 @@ public boolean tryReserve( fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); maxBytesCanReserve = fragment.getMaxBytesCanReserve(reservation, maxBytesCanReserve); long usedBytes = - fragment.sharedBudget + fragment.isSharedBudget() ? fragment.getUsedBytes(reservation.side).get() : reservation.bytesReserved; // Do not publish speculative charges. Two competing requests must not reject one another diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index e78240ec0bce..d7e89cbd6ff3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -293,10 +293,11 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { && pipelineDriverFactories.get(0).getDependencyPipelineIndex() == -1 && (inputExchangeCount == 0 || inputMaxBytes >= maxBlockBytes) && (outputSinkCount == 0 || outputMaxBytes >= maxBlockBytes); + FragmentInstanceContext instanceContext = driverContext.getFragmentInstanceContext(); + instanceContext.initializeExchangeMemoryBudget( + sharedBudget ? inputMaxBytes : 0, sharedBudget ? outputMaxBytes : 0); if (sharedBudget) { - FragmentInstanceId id = getFragmentInstanceId(); - memoryPool.registerFragmentInstanceSharedBudget( - id.getQueryId().getId(), id.getFullId(), inputMaxBytes, outputMaxBytes); + memoryPool.registerFragmentInstanceSharedBudget(instanceContext); } legacySinkLimits.forEach( (sink, legacyLimit) -> diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 136e317e665b..1744216249f8 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -19,8 +19,13 @@ package org.apache.iotdb.db.queryengine.execution.memory; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.commons.memory.MemoryManager; import org.apache.iotdb.db.exception.runtime.MemoryLeakException; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; import com.google.common.util.concurrent.ListenableFuture; @@ -61,6 +66,20 @@ public void before() { QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } + private void registerSharedBudget( + MemoryPool target, String fragmentInstanceId, long inputMaxBytes, long outputMaxBytes) { + FragmentInstanceId id = Mockito.mock(FragmentInstanceId.class); + Mockito.when(id.getQueryId()).thenReturn(new QueryId(QUERY_ID)); + Mockito.when(id.getFullId()).thenReturn(fragmentInstanceId); + FragmentInstanceContext instanceContext = + FragmentInstanceContext.createFragmentInstanceContext( + id, + Mockito.mock(FragmentInstanceStateMachine.class), + Mockito.mock(MemoryReservationManager.class)); + instanceContext.initializeExchangeMemoryBudget(inputMaxBytes, outputMaxBytes); + target.registerFragmentInstanceSharedBudget(instanceContext); + } + @Test public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws Exception { Field reservationsField = MemoryPool.class.getDeclaredField("queryMemoryReservations"); @@ -529,7 +548,7 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { @Test public void testSharedInputBudgetAndIndependentOutput() { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -577,7 +596,7 @@ public void testSideRegistrationCannotChange() { @Test public void testSharedRequestCannotExceedInitializedInputBudget() { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); try { pool.reserveWithPriority( QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, Long.MAX_VALUE, false); @@ -591,7 +610,7 @@ public void testSharedRequestCannotExceedInitializedInputBudget() { @Test public void testHandleLimitsCannotChangeInitializedSharedBudgets() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -626,8 +645,11 @@ public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { new MemoryPool("single-direction", new MemoryManager(1024), 513); for (MemorySide side : MemorySide.values()) { String fi = side.name(); - singleDirectionPool.registerFragmentInstanceSharedBudget( - QUERY_ID, fi, side == MemorySide.INPUT ? 513 : 0, side == MemorySide.OUTPUT ? 513 : 0); + registerSharedBudget( + singleDirectionPool, + fi, + side == MemorySide.INPUT ? 513 : 0, + side == MemorySide.OUTPUT ? 513 : 0); singleDirectionPool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, "active", side); singleDirectionPool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, fi, "unused", side == MemorySide.INPUT ? MemorySide.OUTPUT : MemorySide.INPUT); @@ -649,7 +671,7 @@ public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { @Test public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 128, 128); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 128, 128); ExecutorService executor = Executors.newFixedThreadPool(16); CountDownLatch start = new CountDownLatch(1); CountDownLatch attempted = new CountDownLatch(16); @@ -697,7 +719,7 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { @Test public void testReleaseBeforePendingReservationIsPublished() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); pool = Mockito.spy(pool); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); @@ -750,7 +772,7 @@ public void testReleaseBeforePendingReservationIsPublished() throws Exception { @Test public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -788,7 +810,7 @@ public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Except @Test public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 128, 128); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 128, 128); int count = 16; ExecutorService executor = Executors.newFixedThreadPool(count); CountDownLatch attempted = new CountDownLatch(count); @@ -849,7 +871,7 @@ public void testDefaultFragmentKeepsIndependentPlanNodeQuotas() { @Test public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "a", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "b", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); @@ -873,7 +895,7 @@ public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { @Test public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID, 256, 256); + registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); for (int i = 0; i < 3; i++) { String fi = "other" + i; pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 7d2e23e9d554..21dafba0432e 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -22,6 +22,7 @@ import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.commons.memory.MemoryManager; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; @@ -39,6 +40,7 @@ import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; @@ -58,6 +60,7 @@ import java.util.Collections; import java.util.Comparator; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.doReturn; @@ -89,9 +92,11 @@ public void tearDown() { } private LocalExecutionPlanContext newContext() { - FragmentInstanceContext instance = mock(FragmentInstanceContext.class); - when(instance.getId()) - .thenReturn(new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0")); + FragmentInstanceContext instance = + FragmentInstanceContext.createFragmentInstanceContext( + new FragmentInstanceId(new PlanFragmentId(new QueryId("budget"), 0), "0"), + mock(FragmentInstanceStateMachine.class), + mock(MemoryReservationManager.class)); LocalExecutionPlanContext result = new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); result.addPipelineDriverFactory(mock(Operator.class), result.getDriverContext(), 0); @@ -151,6 +156,15 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side) { } private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit) { + FragmentInstanceContext instance = context.getDriverContext().getFragmentInstanceContext(); + if (shared) { + assertEquals( + limit, + side == MemorySide.INPUT ? instance.getInputMaxBytes() : instance.getOutputMaxBytes()); + } else { + assertEquals(0, instance.getInputMaxBytes()); + assertEquals(0, instance.getOutputMaxBytes()); + } String fi = context.getFragmentInstanceId().getFullId(); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a" + side, side); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b" + side, side); @@ -204,6 +218,7 @@ public void testInputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { context.setExchangeMemoryLimits( node(PlanNodeType.TABLE_COLLECT_NODE, node(PlanNodeType.TABLE_EXCHANGE_NODE)), pool); verify(source).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getOutputMaxBytes()); assertAdmissionIsShared(true); } @@ -222,6 +237,7 @@ public void testOutputOnlyCollectUsesFullBudgetEvenForOneFullBlock() { node(PlanNodeType.SERIES_SCAN))), pool); verify(sink).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getInputMaxBytes()); assertAdmissionIsShared(true, MemorySide.OUTPUT); } @@ -239,6 +255,7 @@ public void testLocalInputsDoNotSplitOutputBudget() { context.setISink(sink); context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); verify(sink).setMaxBytesCanReserve(1024); + assertEquals(0, context.getDriverContext().getFragmentInstanceContext().getInputMaxBytes()); assertAdmissionIsShared(true, MemorySide.OUTPUT); } From 55efd5345ecca9998567829a84f0d8f668414146 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 17:29:36 +0800 Subject: [PATCH 10/15] Simplify FI exchange limits and isolate SQL assertions --- ...2026-10-14-exchange-fanin-budget-design.md | 198 ------------------ .../apache/iotdb/it/env/cluster/EnvUtils.java | 3 +- .../db/it/utils/CollectQueryAssertions.java | 120 +++++++++++ .../apache/iotdb/db/it/utils/TestUtils.java | 46 ++-- .../it/triple/IoTDBShowReceiversTripleIT.java | 84 +++----- .../it/db/it/IoTDBMultiDeviceTableIT.java | 29 ++- .../udf/IoTDBUserDefinedTableFunctionIT.java | 18 +- .../it/insertquery/IoTDBInsertQueryIT.java | 21 +- .../query/recent/IoTDBClusterResultSetIT.java | 91 +++++--- .../it/query/recent/IoTDBCollectFanInIT.java | 12 +- .../it/query/recent/IoTDBCteIT.java | 8 +- .../query/recent/IoTDBTableAggregationIT.java | 5 +- .../it/query/recent/IoTDBTestEnvPortIT.java | 47 ----- .../IoTDBTopKRuntimeFilterScanPathIT.java | 30 ++- .../operator/process/CollectOperator.java | 1 - .../execution/memory/MemoryPool.java | 58 ++--- .../planner/LocalExecutionPlanContext.java | 9 +- .../execution/memory/MemoryPoolTest.java | 92 ++++---- .../LocalExecutionPlanContextBudgetTest.java | 2 + .../commons/client/ClientManagerTest.java | 3 +- 20 files changed, 358 insertions(+), 519 deletions(-) delete mode 100644 docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java delete mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java diff --git a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md b/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md deleted file mode 100644 index 3eb5cbc90625..000000000000 --- a/docs/superpowers/specs/2026-10-14-exchange-fanin-budget-design.md +++ /dev/null @@ -1,198 +0,0 @@ - - -# Exchange 无序扇入与输入输出预算 - -基线：Apache master `aec1ad6139cb4dd1450725f53418cf40c26efb4e`。 -依据：[技术设计文档](https://timechor.feishu.cn/docx/EpDFdmS94ovMCWxluEOc3WGDn8d)，修订 861。 - -## 已确认范围 - -实现仅针对无输入依赖 FI 的输入输出共享账目、Collect/MappingCollect 任意就绪消费，以及由无序汇集引起的计划属性修正。 -聚合使用 HashAggregation：只清空实际输入已失去连续性的 preGroupedSymbols；有效分支内的流式聚合保留。 -MergeSort、Join、ASOF、CrossJoin 不修改消费算法。FI 内只要存在需要特定输入推进的算子，就整组使用原有按 plan node 限额；Collect 下方或上方出现这些算子也不例外。 -TableTopK、CountMerge 的任意就绪消费留到后续；有序输入标志仍需匹配实际输入。 - -## 额度与账目 - -共享采用 FI 级显式启用，默认保留原有配额。整个 FI 必须包含 CollectOperator 或 MappingCollectOperator，且所有多输入汇集都能消费任意就绪输入；顺序消费的 IdentitySink/CountMerge、MergeSort、Join、TopK 及未明确支持的多输入算子都禁止共享。检查在 Exchange 处停止，不跨到其他 FI。 - -在生成全部 pipeline 后再决定预算模式。内部 pipeline 队列具有转发依赖，因此存在多个 pipeline 或调度依赖时整组回退原策略。只含 Collect 的某个子树不会绕过 FI 级否决。 - -通过上述检查后，F 为配置的单 FI 交换预算。按照实际由本 FI 计费的方向分配：只有 INPUT 时 INPUT=F；只有 OUTPUT 时 OUTPUT=F；同时存在两侧时 INPUT=F/2、OUTPUT=F-INPUT；未使用的方向预算为 0。方向在完整本地计划生成后确定，远端 SourceHandle 表示本 FI 的 INPUT，fragment sink 表示 OUTPUT；LocalSourceHandle 的队列由生产 FI 计费，不重复占用接收 FI 的 INPUT。仅实际启用的方向须容纳配置的最大 TsBlock，否则回退原配额。未启用共享的 FI 恢复原有 exchangeSumNum 配额计算和按 plan node 的申请上限；不能只恢复限额数值却仍按同侧总量进行准入。 -远端 SourceHandle 计入接收 FI 的 INPUT；SinkChannel 和 SharedTsBlockQueue 计入生产 FI 的 OUTPUT。 -本地共享队列只计费一次，创建先后不改变所有者。 - -MemoryPool 维护 query→FI→planNode 明细，注册时固定 side。 -本地计划生成全部算子和 pipeline 后，一次性检查共享资格并计算两个方向的预算，写入 `FragmentInstanceContext` 的 `inputMaxBytes`、`outputMaxBytes` 字段。`MemorySide` 仅表示方向，不计算预算，也不接受 `hasInput`/`hasOutput` 参数。MemoryPool 的 FI 账目仅保留该上下文的引用，不复制预算字段。共享模式的申请、单块上限检查和等待唤醒统一读取 FI 保存的预算；handle 传入的限额仅用于原有独立配额模式。回退独立配额时上下文的两个共享预算字段为零。 -FI 两侧保留 AtomicLong 统计，plan node 保留字节明细。申请和释放在 FI 的记账锁内更新；全局池采用有界原子申请。申请前先检查适用的同侧总量或 plan node 限额，失败不产生临时计费，也不需要回滚，避免并发请求分别撞到不同限额后全部挂起。 - -释放后在 FI 记账锁外完成等待 Future；保留入队后的重试和每个 Future 的完成/取消同步。稳定快照中的 FI 同侧总量等于同侧明细之和。 -保留现有最高优先级查询的零计费兜底与交换组件的收发、等待、释放机制。 - -## Collect 生命周期 - -AnyChildBlocked 缓存直接 child.isBlocked 的 Future，每个 pending 周期只有一个 listener。 -新轮次重定向现有 listener；发布后检查完成状态，覆盖注册与完成交错。 -回调只记录就绪或错误，在锁外完成选择 Future；所有 child 方法由 Driver 线程调用。 -重复 isBlocked 与输出背压保留选择；取消对外等待不向 child 传播。 - -一次 next 消耗一次选择。null/空块刷新等待，不表示结束。 -hasNext=false 后检查 isFinished，确认结束才关闭 child 并移除记录。 -MappingCollect 按实际 childId 映射。关闭唤醒待完成选择，忽略迟到回调，逐一尝试关闭全部输入。 -保留内存 R 为各存活 child 保留内存之和；峰值为 max(R,max(R-r_i+p_i))；返回大小为 max(b_i)。 - -## 计划属性 - -多路 Collect/Union 打断分支连续性和排序；单路可透传。 -Filter 与符号透传 Project 不恢复属性；有效 Sort/MergeSort 可以重新建立属性。 -在 SINGLE/PARTIAL/FINAL 实际 child 组装后校验聚合，包括 DISTINCT/mask、普通分阶段聚合及 Union 下推。 -Group 在分区连续性失效时改为完整 Sort。复用 master 的 StreamSort 前缀校验和 Window/Fill 修复。 -SortElimination 在各汇集分支独立优化，并阻止跨无序汇集删除恢复排序。 - -## 验证 - -覆盖纯 Collect/MappingCollect 显式启用、混合 FI 全组回退、Exchange 边界、内部 pipeline 否决、单块预算边界，以及共享/原配额 FI 共存； -覆盖共享额度并发、INPUT/OUTPUT 隔离、所有者、奇数 F、超大申请、失败准入无计费、释放取消及稳定账目一致性； -覆盖 100/1000 child、A 阻塞 B 就绪、listener 复用、注册竞争、输出背压、空块、结束、异常取消、关闭和映射； -覆盖文档 SQL 对应的聚合、Group/Window、StreamSort、Union 与 Fill 计划及实际结果。 -运行 Spotless、英文/中文全 reactor test-compile、针对性单测，以及本分支新增或修改的 IT。 - -## 实现验证记录（2026-10-14） - -上述约定范围的实现已完成。单元测试 86 项通过；IoTDBCollectFanInIT 的 6 项测试在单节点和三 DataNode 集群中分别通过。 -集成验证包含普通与 DISTINCT 聚合、Union 分阶段聚合、嵌套窗口、Fill、排序/TopK、Join 依赖、树模型以及树表视图。 - -构建先执行以下干净构建，生成当前分支的节点包并安装本地依赖： - -```sh -mvn -o clean install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests -``` - -全 reactor 中文编译与英文编译/回归均通过；英文 test 阶段包含完整 test-compile： - -```sh -mvn -o test-compile -DskipTests -P with-zh-locale -mvn -o test -Dtest=CollectOperatorTest,MemoryPoolTest,LocalExecutionPlanContextBudgetTest,CollectInputPropertyTest,SharedTsBlockQueueTest,SourceHandleTest,SinkChannelTest,LocalSourceHandleTest,LocalSinkChannelTest,MPPDataExchangeManagerTest,ParallelizeGroupingTest,TableFunctionSortTest -Dsurefire.failIfNoSpecifiedTests=false -``` - -两个集成环境均只运行本分支新增的测试类： - -```sh -mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableSimpleIT -P with-integration-tests -mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableClusterIT -P with-integration-tests -``` - -使用离线模式是为了避免远端 SNAPSHOT ZIP 下载停滞；节点执行的是当前分支重新打包的代码。 -Spotless、Checkstyle、新增文件许可证及中英文消息键名/占位符一致性检查通过。 - -## 自查与补测（2026-10-14） - -自查复现了 MemoryPool 的等待申请漏唤醒：一次申请额度不足后，旧申请先释放全部额度， -新申请才进入等待队列；释放线程看不到新申请，导致空闲额度下仍无限等待。 -回归测试用 latch 固定这一时序，在修复前稳定失败。 -修复在申请入队后再次检查额度，并与 free/tryCancel 共用该 future 的锁，避免重复计费。 - -新增 8 项单元测试：等待申请发布竞态、取消不占用共享额度、16 路阻塞申请反复获得额度后账目归零； -空 Collect、全部阻塞时异常唤醒、并发就绪后每路恰好消费一次； -Collect 后 TopK 不误用有序输入，以及有效分支保留流式聚合。 -增强关闭异常测试，检查其他输入仍被关闭且后续异常被保留。 -新增一项集成测试，验证全局窗口后按字段倒序的 TopK 跨设备结果及完整窗口计数。 - -补测后的 118 项定向单测通过（14 个测试类，零失败、零跳过）；全 reactor 中英文编译通过。 -在前述英文回归命令中加入 `OrderingPropertyPropagationTest,TopKRuntimeFilterDistributedPlanTest`， -同时验证有效排序优化与 TopK runtime filter 不受影响。 -重新执行 `mvn -o install -DskipTests -pl integration-test,iotdb-core/confignode,iotdb-client/cli -am -P with-integration-tests` -后运行前述两条 verify 命令；IoTDBCollectFanInIT 的 7 项测试在单节点和三 DataNode 集群中分别全部通过。 -Spotless、Checkstyle、许可证与 `git diff --check` 检查通过。 - -## FI 共享范围修正验证（2026-10-15） - -共享改为完整 FI 的显式启用：仅支持 Collect/MappingCollect 的任意就绪汇集； -同一 FI 内的多输入依赖算子、内部 pipeline 队列或调度依赖会使整组回退。 -Exchange 边界之外的上游 FI 不参与本 FI 的判断。回退同时恢复 handle 限额和按 plan node 准入。 - -107 项定向单元测试通过，包括纯 Collect/Union、混合 FI、pipeline 依赖、单块额度边界、 -共享与独立配额共存、并发准入，以及真实 TableMergeSort 配合模拟输入和 MemoryPool 的回归。 -本地双队列转发回归也通过。Spotless、Checkstyle、许可证和 `git diff --check` 检查通过。 - -```sh -mvn -o test -pl iotdb-core/datanode -Dtest=LocalExecutionPlanContextBudgetTest,MemoryPoolTest,SharedTsBlockQueueTest,SourceHandleTest,SinkChannelTest,LocalSourceHandleTest,LocalSinkChannelTest,MPPDataExchangeManagerTest,PipelineBuilderTest,CollectInputPropertyTest -mvn -o verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -am -P TableSimpleIT -P with-integration-tests -mvn -o failsafe:integration-test@integration-test failsafe:verify@verify -DskipUTs -Dit.test=IoTDBCollectFanInIT -DfailIfNoTests=false -Dfailsafe.failIfNoSpecifiedTests=false -pl integration-test -P TableClusterIT -P with-integration-tests -``` - -单节点 verify 从当前源码构建依赖和测试节点；集群运行复用同一份节点模板。 -IoTDBCollectFanInIT 的 8 项测试在两种环境中均通过，零失败、零跳过。 - -提交前再次执行完整 reactor 的中英文编译，两者均通过： - -```sh -mvn -o test-compile -DskipTests -P with-zh-locale -mvn -o test-compile -DskipTests -``` - -本次生产代码没有新增未本地化的字符串，且没有修改中英文消息常量。 - -## 单方向预算修正验证（2026-10-15） - -该轮修正取消默认按两侧均分，改为根据本 FI 实际计费方向分配： -单方向使用完整 F，双方向才拆分，未使用方向为 0。本次已将计算统一收回本地计划初始化， -不再保留该轮引入的 `MemorySide.getMaxBytes(F, hasInput, hasOutput)` 接口。 -最大 TsBlock 校验只作用于启用的方向；原有 Collect 资格和依赖回退规则保持不变。 - -新增输入单向、输出单向、B=F、本地输入和本地/远端混合输入用例； -同时验证奇数 F、未使用方向及每个单方向一次申请完整 F 的边界。 -113 项定向单测全部通过。首轮中单块边界测试的默认测试池上限为 512、申请为 513， -将该测试池配置为目标 F=513 后重跑 MemoryPoolTest，38 项全部通过，其余 75 项首轮通过。 - -完整 reactor 中英文 test-compile、Spotless、Checkstyle 和许可证检查通过。 -重新从当前代码打包节点后，IoTDBCollectFanInIT 在 TableSimpleIT 和三 DataNode 的 -TableClusterIT 中各 8 项全部通过；集群复用同一份节点模板，节点包哈希与本次构建一致。 - -## FI 预算字段初始化验证（2026-10-15） - -`FragmentMemoryReservation` 保存 `inputMaxBytes`、`outputMaxBytes`； -完整本地计划生成后按输入 exchange、fragment sink 和依赖检查结果一次性初始化。 -共享申请、单块上限检查与等待重试均读取 FI 字段，独立配额模式仍读取 handle 限额。 -`MemorySide` 只保留方向枚举，移除布尔参数预算接口。 - -115 项定向单测通过，覆盖 handle 参数不能增减共享上限、超额请求不进入等待队列、 -未使用方向预算为零、奇数预算拆分和原有独立配额回退。 -完整 reactor 中英文编译、Spotless、Checkstyle 和许可证检查通过；生产代码没有新增字符串字面量。 - -重新打包当前源码后，`IoTDBCollectFanInIT` 在 TableSimpleIT 和三 DataNode 的 TableClusterIT -中各 8 项全部通过，无失败或跳过。集群复用同一份节点模板，其服务端 jar 哈希与本次构建一致。 - -## FI 预算配置归属修正验证（2026-10-15） - -预算字段移至 `FragmentInstanceContext`，由完整本地计划生成后的初始化流程赋值。 -MemoryPool 的 FI 账目保存上下文引用，通过 getter 读取限额，不再复制预算或保存共享布尔字段。 -该引用随 FI 记账记录注销；预算仍在 driver 启动前初始化。 -独立配额回退时上下文的两个共享预算字段为零，不注册共享预算。 - -123 项定向单测通过，包括真实 FI 上下文的初始化字段断言、申请上限和等待唤醒回归， -以及已有的 FI 上下文生命周期测试。生产代码没有新增字符串字面量。 - -本次单测在前述 `-Dtest` 列表中加入 `FragmentInstanceContextTest`。完整 reactor 的中英文 -`test-compile -DskipTests`、Spotless、Checkstyle、许可证与 `git diff --check` 检查均通过。 -重新打包节点后，`IoTDBCollectFanInIT` 在 TableSimpleIT 和三 DataNode 的 TableClusterIT 中各 -8 项全部通过，无失败或跳过；测试模板与当前服务端 jar 的哈希一致。 diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java index e6c58e251f59..b69e86c12bd4 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/EnvUtils.java @@ -180,8 +180,7 @@ public static Map listPortOccupationWindows(final List p */ public static Map listPortOccupationUnix(final List ports) throws IOException { - // The last column is the socket state (e.g. LISTEN); the address precedes it. - return listPortOccupation(ports, "lsof -iTCP -sTCP:LISTEN,TIME_WAIT -P -n", 10, 8, 1); + return listPortOccupation(ports, "lsof -iTCP -sTCP:LISTEN,TIME_WAIT -P -n", 10, 9, 1); } private static Pair getClusterNodesNum(final int index) { diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java new file mode 100644 index 000000000000..e53ad150ca43 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/CollectQueryAssertions.java @@ -0,0 +1,120 @@ +/* + * 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.it.utils; + +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QueryBody; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.QuerySpecification; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.TableSubquery; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Explain; +import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser; +import org.apache.iotdb.itbase.runtime.ClusterTestResultSet; + +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/** Explicit SQL-aware assertions for queries affected by ready-input Collect consumption. */ +public final class CollectQueryAssertions { + private CollectQueryAssertions() {} + + /** Only queries without final ordering or an arbitrary top-level slice may use this path. */ + public static List getUnorderedResultSets(String sql, ResultSet resultSet) { + org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement statement = + new SqlParser().createStatement(sql, ZoneOffset.UTC, null); + assertTrue("Unordered assertions require a SELECT query: " + sql, statement instanceof Query); + Query query = (Query) statement; + while (true) { + assertFalse( + "Final ORDER BY requires ordered assertions: " + sql, query.getOrderBy().isPresent()); + assertFalse( + "LIMIT requires deterministic result selection: " + sql, query.getLimit().isPresent()); + assertFalse( + "OFFSET requires deterministic result selection: " + sql, query.getOffset().isPresent()); + QueryBody body = query.getQueryBody(); + if (body instanceof QuerySpecification) { + QuerySpecification specification = (QuerySpecification) body; + assertFalse( + "Final ORDER BY requires ordered assertions: " + sql, + specification.getOrderBy().isPresent()); + assertFalse( + "LIMIT requires deterministic result selection: " + sql, + specification.getLimit().isPresent()); + assertFalse( + "OFFSET requires deterministic result selection: " + sql, + specification.getOffset().isPresent()); + } + if (!(body instanceof TableSubquery)) { + break; + } + // Parentheses around the whole query retain its ordering; nested FROM/UNION branches do not. + query = ((TableSubquery) body).getQuery(); + } + return nodeResults(resultSet); + } + + /** EXPLAIN text may contain node-specific details; each node must satisfy the plan assertions. */ + public static List getExplainResultSets(String sql, ResultSet resultSet) { + assertTrue( + "Plan assertions require EXPLAIN: " + sql, + new SqlParser().createStatement(sql, ZoneOffset.UTC, null) instanceof Explain); + return nodeResults(resultSet); + } + + private static List nodeResults(ResultSet resultSet) { + return resultSet instanceof ClusterTestResultSet + ? ClusterTestResultSet.getNodeResultSets(resultSet) + : Collections.singletonList(resultSet); + } + + /** Compare every node independently while retaining duplicate multiplicity. */ + public static void assertUnorderedQueryResult( + String sql, ResultSet resultSet, String expectedHeader, List expectedRows) + throws SQLException { + List sortedExpected = new ArrayList<>(expectedRows); + Collections.sort(sortedExpected); + for (ResultSet nodeResult : getUnorderedResultSets(sql, resultSet)) { + ResultSetMetaData metadata = nodeResult.getMetaData(); + StringBuilder header = new StringBuilder(); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + header.append(metadata.getColumnName(i)).append(","); + } + assertEquals(expectedHeader, header.toString()); + List actualRows = new ArrayList<>(); + while (nodeResult.next()) { + StringBuilder row = new StringBuilder(); + for (int i = 1; i <= metadata.getColumnCount(); i++) { + row.append(nodeResult.getString(i)).append(","); + } + actualRows.add(row.toString()); + } + Collections.sort(actualRows); + assertEquals(sortedExpected, actualRows); + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java index a85aacc4c368..bf1cb4376bb9 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java @@ -63,7 +63,6 @@ import static org.apache.iotdb.itbase.constant.TestConstant.TIMESTAMP_STR; import static org.apache.iotdb.itbase.env.BaseEnv.TABLE_SQL_DIALECT; import static org.apache.iotdb.itbase.env.BaseEnv.TREE_SQL_DIALECT; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.awaitility.Awaitility.await; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -754,15 +753,6 @@ public static void assertResultSetEqual( final String expectedHeader, final Collection expectedResult, final boolean printResult) { - if (expectedResult instanceof Set) { - List nodeResults = getNodeResultSets(actualResultSet); - if (nodeResults.size() != 1 || nodeResults.get(0) != actualResultSet) { - for (ResultSet nodeResult : nodeResults) { - assertResultSetEqual(nodeResult, expectedHeader, expectedResult, printResult); - } - return; - } - } try { final ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData(); final StringBuilder header = new StringBuilder(); @@ -824,31 +814,27 @@ public static void assertResultSetEqual( String expectedHeader, Set expectedRetSet, Consumer consumer) { - assertResultSetEqual(actualResultSet, expectedHeader, expectedRetSet, false); - } - - /** Checks unordered rows on every node, preserving duplicate multiplicity. */ - public static void assertResultSetEqualUnordered( - ResultSet resultSet, String expectedHeader, List expectedRows) throws SQLException { - List sortedExpected = new ArrayList<>(expectedRows); - Collections.sort(sortedExpected); - for (ResultSet nodeResult : getNodeResultSets(resultSet)) { - ResultSetMetaData metadata = nodeResult.getMetaData(); + try { + ResultSetMetaData resultSetMetaData = actualResultSet.getMetaData(); StringBuilder header = new StringBuilder(); - for (int i = 1; i <= metadata.getColumnCount(); i++) { - header.append(metadata.getColumnName(i)).append(","); + for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) { + header.append(resultSetMetaData.getColumnName(i)).append(","); } assertEquals(expectedHeader, header.toString()); - List actualRows = new ArrayList<>(); - while (nodeResult.next()) { - StringBuilder row = new StringBuilder(); - for (int i = 1; i <= metadata.getColumnCount(); i++) { - row.append(nodeResult.getString(i)).append(","); + + Set actualRetSet = new HashSet<>(); + + while (actualResultSet.next()) { + StringBuilder builder = new StringBuilder(); + for (int i = 1; i <= resultSetMetaData.getColumnCount(); i++) { + builder.append(actualResultSet.getString(i)).append(","); } - actualRows.add(row.toString()); + actualRetSet.add(builder.toString()); } - Collections.sort(actualRows); - assertEquals(sortedExpected, actualRows); + assertStringSetEqual(expectedRetSet, actualRetSet); + } catch (Exception e) { + e.printStackTrace(); + Assert.fail(String.valueOf(e)); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java index fda9935e3191..0712387775b9 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java @@ -44,7 +44,6 @@ import java.sql.Connection; import java.sql.ResultSet; -import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.Arrays; @@ -55,8 +54,6 @@ import java.util.Set; import java.util.concurrent.TimeUnit; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; - @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT3.class}) public class IoTDBShowReceiversTripleIT { @@ -217,44 +214,39 @@ private boolean hasReceiversOnDifferentDataNodes( final String secondPipeName, final int secondReceiverDataNodeId) throws SQLException { + boolean hasFirstPipe = false; + boolean hasSecondPipe = false; + final Set receiverDataNodeIds = new HashSet<>(); try (final Connection connection = receiverEnv.getConnection( SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); final Statement statement = connection.createStatement(); - final ResultSet clusterResult = statement.executeQuery(sql)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { - boolean hasFirstPipe = false; - boolean hasSecondPipe = false; - final Set receiverDataNodeIds = new HashSet<>(); - while (resultSet.next()) { - final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); - if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) - || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) - || pipeIds == null) { - continue; - } - final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); - if (receiverDataNodeId == firstReceiverDataNodeId - && pipeIds.contains(firstPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasFirstPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); - } - if (receiverDataNodeId == secondReceiverDataNodeId - && pipeIds.contains(secondPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasSecondPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); - } + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) + || pipeIds == null) { + continue; + } + final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); + if (receiverDataNodeId == firstReceiverDataNodeId + && pipeIds.contains(firstPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasFirstPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); } - if (!hasFirstPipe || !hasSecondPipe || receiverDataNodeIds.size() < 2) { - return false; + if (receiverDataNodeId == secondReceiverDataNodeId + && pipeIds.contains(secondPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasSecondPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); } } } - return true; + return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2; } private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception { @@ -276,28 +268,20 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception private static String getString( final ResultSet resultSet, final String treeColumnName, final String tableColumnName) throws SQLException { - return resultSet.getString(getColumnIndex(resultSet, treeColumnName, tableColumnName)); + try { + return resultSet.getString(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getString(tableColumnName); + } } private static int getInt( final ResultSet resultSet, final String treeColumnName, final String tableColumnName) throws SQLException { - return resultSet.getInt(getColumnIndex(resultSet, treeColumnName, tableColumnName)); - } - - private static int getColumnIndex( - final ResultSet resultSet, final String treeColumnName, final String tableColumnName) - throws SQLException { - // SHOW RECEIVERS and information_schema.receivers expose different column names. - // Resolve from metadata because JDBC getters do not all use SQLException for unknown names. - final ResultSetMetaData metadata = resultSet.getMetaData(); - for (int column = 1; column <= metadata.getColumnCount(); column++) { - final String name = metadata.getColumnName(column); - if (treeColumnName.equals(name) || tableColumnName.equals(name)) { - return column; - } + try { + return resultSet.getInt(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getInt(tableColumnName); } - throw new AssertionError( - "Missing receiver column: " + treeColumnName + " / " + tableColumnName); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java index 096c1383c10d..999f16e4b080 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBMultiDeviceTableIT.java @@ -37,7 +37,6 @@ import java.util.HashMap; import java.util.Map; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -265,27 +264,25 @@ public void testMultiDeviceQueryAndDelete() { } private void testSelectAll() { - String selectSql = "select * from t"; + String selectSql = "select * from t ORDER BY tag1, tag2, time"; try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT); Statement statement = connection.createStatement()) { statement.execute("use \"test\""); - try (ResultSet clusterResult = statement.executeQuery(selectSql)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { - Map lastTimeMap = new HashMap<>(); - int cnt = 0; - while (resultSet.next()) { - String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2"); - long before = lastTimeMap.getOrDefault(tag, -1L); - long cur = resultSet.getTimestamp("time").getTime(); - if (cur <= before) { - fail("time order wrong!"); - } - lastTimeMap.put(tag, cur); - cnt++; + Map lastTimeMap = new HashMap<>(); + try (ResultSet resultSet = statement.executeQuery(selectSql)) { + int cnt = 0; + while (resultSet.next()) { + String tag = resultSet.getString("tag1") + "_" + resultSet.getString("tag2"); + long before = lastTimeMap.getOrDefault(tag, -1L); + long cur = resultSet.getTimestamp("time").getTime(); + if (cur <= before) { + fail("time order wrong!"); } - assertEquals(16030, cnt); + lastTimeMap.put(tag, cur); + cnt++; } + assertEquals(16030, cnt); } } catch (Exception e) { e.printStackTrace(); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java index 25967b8caeb8..25c7176ca9fc 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/udf/IoTDBUserDefinedTableFunctionIT.java @@ -39,9 +39,9 @@ import java.util.HashSet; import java.util.Set; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.fail; @RunWith(IoTDBTestRunner.class) @@ -204,14 +204,14 @@ public void testLargeResultIsSplitWithoutDataLoss() throws Exception { try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); - try (ResultSet clusterResult = - statement.executeQuery( - "SELECT * FROM large_result(vehicle, " - + LARGE_RESULT_REPEAT_COUNT - + ", " - + LARGE_RESULT_PAYLOAD_SIZE - + ")")) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + String sql = + "SELECT * FROM large_result(vehicle, " + + LARGE_RESULT_REPEAT_COUNT + + ", " + + LARGE_RESULT_PAYLOAD_SIZE + + ")"; + try (ResultSet clusterResult = statement.executeQuery(sql)) { + for (ResultSet resultSet : getUnorderedResultSets(sql, clusterResult)) { Set returnedRows = new HashSet<>(); while (resultSet.next()) { int repeatIndex = resultSet.getInt("repeat_index"); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java index e56aadf5c86a..6c3ce543300f 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java @@ -48,7 +48,6 @@ import java.util.List; import java.util.Locale; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; @@ -280,17 +279,15 @@ public void testPartialColumn() throws SQLException { statement.execute("DELETE FROM vehicle1"); statement.execute( "INSERT INTO vehicle1(time, s2, manufacturer) SELECT time, s2, manufacturer FROM vehicle0"); - try (ResultSet clusterResult = - statement.executeQuery("SELECT deviceId, manufacturer, s1, s2 FROM vehicle1")) { - for (ResultSet nodeResult : getNodeResultSets(clusterResult)) { - while (nodeResult.next()) { - for (int i = 1; i <= 4; i++) { - if (i % 2 == 1) { - assertNull(nodeResult.getString(i)); - } else { - assertNotNull(nodeResult.getString(i)); - } - } + resultSet = + statement.executeQuery( + "SELECT deviceId, manufacturer, s1, s2 FROM vehicle1 ORDER BY manufacturer, s2"); + while (resultSet.next()) { + for (int i = 1; i <= 4; i++) { + if (i % 2 == 1) { + assertNull(resultSet.getString(i)); + } else { + assertNotNull(resultSet.getString(i)); } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java index a29c9793b72d..332c69477bda 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBClusterResultSetIT.java @@ -38,9 +38,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.Set; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; @@ -49,48 +50,49 @@ public class IoTDBClusterResultSetIT { @Test - public void testUnorderedSetChecksEveryNode() throws SQLException { + public void testUnorderedRowsCheckEveryNode() throws SQLException { try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { - TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,")); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "b,")); } try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "c"})) { - assertThrows( - AssertionError.class, - () -> TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"))); - } - } - - @Test - public void testUnorderedSetConsumerOverloadChecksEveryNode() throws SQLException { - try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { - TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"), message -> {}); - } - try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"a"})) { assertThrows( AssertionError.class, () -> - TestUtils.assertResultSetEqual(result, "value,", Set.of("a,", "b,"), message -> {})); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "b,"))); } } @Test public void testUnorderedRowsPreserveDuplicateMultiplicity() throws SQLException { try (ResultSet result = cluster(new String[] {"a", "a", "b"}, new String[] {"b", "a", "a"})) { - TestUtils.assertResultSetEqualUnordered(result, "value,", Arrays.asList("a,", "a,", "b,")); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "a,", "b,")); } - for (String[] rows : new String[][] {{"a", "b"}, {"a", "a", "a", "b"}}) { - try (ResultSet result = cluster(new String[] {"a", "a", "b"}, rows)) { + for (String[] values : new String[][] {{"a", "b"}, {"a", "a", "a", "b"}}) { + try (ResultSet result = cluster(new String[] {"a", "a", "b"}, values)) { assertThrows( AssertionError.class, () -> - TestUtils.assertResultSetEqualUnordered( - result, "value,", Arrays.asList("a,", "a,", "b,"))); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("a,", "a,", "b,"))); } } } @Test - public void testOrderedRowsStillRejectDifferentNodeOrder() throws SQLException { + public void testFinalOrderingRejectsUnorderedAssertions() throws SQLException { + for (String sql : + new String[] { + "SELECT value FROM t ORDER BY value", + "(SELECT value FROM t ORDER BY value)", + "SELECT value FROM t UNION ALL SELECT value FROM t ORDER BY value" + }) { + try (ResultSet result = rows("a")) { + assertThrows(AssertionError.class, () -> getUnorderedResultSets(sql, result)); + } + } try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { assertThrows( AssertionError.class, @@ -98,15 +100,50 @@ public void testOrderedRowsStillRejectDifferentNodeOrder() throws SQLException { } } + @Test + public void testNestedOrderingDoesNotOrderFinalResult() throws SQLException { + for (String sql : + new String[] { + "(SELECT value FROM t ORDER BY value LIMIT 1) UNION ALL (SELECT value FROM t ORDER BY value LIMIT 1)", + "SELECT * FROM (SELECT value FROM t ORDER BY value)", + "SELECT row_number() OVER (ORDER BY value) FROM t", + "SELECT 'order by' FROM t /* ORDER BY value */" + }) { + try (ResultSet result = cluster(new String[] {"a", "b"}, new String[] {"b", "a"})) { + assertUnorderedQueryResult(sql, result, "value,", Arrays.asList("a,", "b,")); + } + } + } + + @Test + public void testArbitrarySlicesNeedDeterministicSelection() throws SQLException { + for (String sql : + new String[] {"SELECT value FROM t LIMIT 1", "SELECT value FROM t OFFSET 1"}) { + try (ResultSet result = rows("a")) { + assertThrows(AssertionError.class, () -> getUnorderedResultSets(sql, result)); + } + } + } + @Test public void testSingleNodeAndNullValues() throws SQLException { try (ResultSet result = rows("a", null, "a")) { - assertEquals(1, getNodeResultSets(result).size()); - TestUtils.assertResultSetEqualUnordered(result, "value,", Arrays.asList("null,", "a,", "a,")); + assertEquals(1, getUnorderedResultSets("SELECT value FROM t", result).size()); + assertUnorderedQueryResult( + "SELECT value FROM t", result, "value,", Arrays.asList("null,", "a,", "a,")); } try (ResultSet result = cluster(new String[] {"a"})) { - assertEquals(1, getNodeResultSets(result).size()); - TestUtils.assertResultSetEqual(result, "value,", Set.of("a,")); + assertEquals(1, getUnorderedResultSets("SELECT value FROM t", result).size()); + assertUnorderedQueryResult("SELECT value FROM t", result, "value,", Arrays.asList("a,")); + } + } + + @Test + public void testExplainChecksAreExplicit() throws SQLException { + try (ResultSet result = cluster(new String[] {"plan-a"}, new String[] {"plan-b"})) { + assertThrows(AssertionError.class, () -> getExplainResultSets("SELECT value FROM t", result)); + assertEquals( + 2, getExplainResultSets("EXPLAIN SELECT value FROM t ORDER BY value", result).size()); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java index f9de8a0664fb..943a2150d7f6 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -36,7 +36,7 @@ import java.util.ArrayList; import java.util.List; -import static org.apache.iotdb.db.it.utils.TestUtils.assertResultSetEqualUnordered; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -99,11 +99,11 @@ public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Excepti try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE); - try (ResultSet result = - statement.executeQuery( - "SELECT device, zone, value FROM readings1 UNION ALL " - + "SELECT device, zone, value FROM readings1")) { - assertResultSetEqualUnordered(result, "device,zone,value,", expectedRows); + String sql = + "SELECT device, zone, value FROM readings1 UNION ALL " + + "SELECT device, zone, value FROM readings1"; + try (ResultSet result = statement.executeQuery(sql)) { + assertUnorderedQueryResult(sql, result, "device,zone,value,", expectedRows); } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java index 87571cfb05db..7ada8c567bd6 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCteIT.java @@ -57,9 +57,9 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -364,7 +364,7 @@ public void testNestExplain1() throws SQLException { // explain try (ResultSet clusterResult = statement.executeQuery(sql)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { ResultSetMetaData metaData = resultSet.getMetaData(); assertEquals(metaData.getColumnCount(), 1); assertEquals(metaData.getColumnName(1), "distribution plan"); @@ -394,7 +394,7 @@ public void testNestExplain2() throws SQLException { // explain try (ResultSet clusterResult = statement.executeQuery(sql)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { ResultSetMetaData metaData = resultSet.getMetaData(); assertEquals(metaData.getColumnCount(), 1); assertEquals(metaData.getColumnName(1), "distribution plan"); @@ -424,7 +424,7 @@ public void testExplainJoin() throws SQLException { // explain try (ResultSet clusterResult = statement.executeQuery(sql)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + for (ResultSet resultSet : getExplainResultSets(sql, clusterResult)) { ResultSetMetaData metaData = resultSet.getMetaData(); assertEquals(metaData.getColumnCount(), 1); assertEquals(metaData.getColumnName(1), "distribution plan"); diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java index 99188903b99a..d707cd65bdc5 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTableAggregationIT.java @@ -41,10 +41,10 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getExplainResultSets; import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData; import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; -import static org.apache.iotdb.itbase.runtime.ClusterTestResultSet.getNodeResultSets; import static org.apache.iotdb.relational.it.db.it.IoTDBMultiTAGsWithAttributesTableIT.buildHeaders; import static org.apache.iotdb.relational.it.db.it.IoTDBMultiTAGsWithAttributesTableIT.repeatTest; @@ -5904,7 +5904,8 @@ private void assertRateQueryUsesIntermediateAggregation(String query) { Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE_NAME); try (ResultSet clusterResult = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) { - for (ResultSet resultSet : getNodeResultSets(clusterResult)) { + for (ResultSet resultSet : + getExplainResultSets("EXPLAIN (FORMAT JSON) " + query, clusterResult)) { Assert.assertTrue(resultSet.next()); String plan = resultSet.getString(1); // The PARTIAL/FINAL split only exists when the scanned data spans multiple DataRegions. diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java deleted file mode 100644 index dbb9a9d7a74c..000000000000 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTestEnvPortIT.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * 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.relational.it.query.recent; - -import org.apache.iotdb.it.env.cluster.EnvUtils; -import org.apache.iotdb.itbase.category.TableClusterIT; -import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; - -import org.junit.Test; -import org.junit.experimental.categories.Category; - -import java.net.InetAddress; -import java.net.ServerSocket; -import java.util.Collections; -import java.util.Map; - -import static org.junit.Assert.assertEquals; - -@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) -public class IoTDBTestEnvPortIT { - - @Test - public void testListeningPortIsReportedAsOccupied() throws Exception { - try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getByName("127.0.0.1"))) { - Map occupiedPorts = - EnvUtils.listPortOccupation(Collections.singletonList(server.getLocalPort())); - assertEquals( - Long.valueOf(ProcessHandle.current().pid()), occupiedPorts.get(server.getLocalPort())); - } - } -} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java index a0a638dcfe0e..59499aa5db0d 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBTopKRuntimeFilterScanPathIT.java @@ -41,14 +41,11 @@ import java.io.IOException; import java.nio.file.Files; import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import static org.apache.iotdb.db.it.utils.TestUtils.assertResultSetEqualUnordered; import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData; import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest; @@ -134,23 +131,20 @@ public void orderByTimeAscLimitOnTable() { } @Test - public void unionAllSiblingTopKOrderByTimeLimit() throws SQLException { - // Each branch has an ORDER BY, but UNION ALL does not order the combined result. - String sql = + public void unionAllSiblingTopKOrderByTimeLimit() { + String[] expectedHeader = new String[] {"time", "device_id", "s1"}; + String[] retArray = + new String[] { + "1970-01-01T00:00:00.030Z,d1,111,", "1970-01-01T00:00:00.032Z,d3,333,", + }; + tableResultSetEqualTest( "(SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd1' ORDER BY time DESC LIMIT 1)" + " UNION ALL" - + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)"; - try (Connection connection = EnvFactory.getEnv().getTableConnection(); - Statement statement = connection.createStatement()) { - connection.setClientInfo("time_zone", "+00:00"); - statement.execute("USE " + TABLE_DATABASE); - try (ResultSet resultSet = statement.executeQuery(sql)) { - assertResultSetEqualUnordered( - resultSet, - "time,device_id,s1,", - Arrays.asList("1970-01-01T00:00:00.030Z,d1,111,", "1970-01-01T00:00:00.032Z,d3,333,")); - } - } + + " (SELECT time, device_id, s1 FROM table1 WHERE device_id = 'd3' ORDER BY time DESC LIMIT 1)" + + " ORDER BY time, device_id, s1", + expectedHeader, + retArray, + TABLE_DATABASE); } @Test diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java index 2cc5deec59ea..bd7dec028b13 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java @@ -208,7 +208,6 @@ public long ramBytesUsed() { return INSTANCE_SIZE + anyChildBlocked.ramBytesUsed() + children.stream() - .filter(java.util.Objects::nonNull) .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject) .sum() + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 8c3ebcbbc9fb..219372c9bd34 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -25,7 +25,6 @@ import org.apache.iotdb.commons.utils.TestOnly; import org.apache.iotdb.db.exception.runtime.MemoryLeakException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; -import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import com.google.common.util.concurrent.AbstractFuture; import com.google.common.util.concurrent.ListenableFuture; @@ -60,8 +59,8 @@ public static class MemoryReservationFuture extends AbstractFuture { /** * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the - * pool. This is the per-plan-node limit for legacy fragments. Shared fragments always use their - * initialized FI-side limit, including when retrying a pending reservation. + * pool. Local planning gives handles their FI-side limit in shared mode or their independent + * plan-node limit in legacy mode. Retain the supplied limit when retrying this request. */ private final long maxBytesCanReserve; @@ -155,7 +154,6 @@ public long getReservedBytes() { private final String id; private final IMemoryBlock memoryBlock; - private final long maxBytesPerFragmentInstance; private static class PlanNodeReservation { private final MemorySide side; @@ -175,20 +173,8 @@ private static class FragmentMemoryReservation { private final Map planNodes = new ConcurrentHashMap<>(); private final AtomicLong inputUsedBytes = new AtomicLong(); private final AtomicLong outputUsedBytes = new AtomicLong(); - // Registered after local planning; the FI context owns the initialized shared limits. - private FragmentInstanceContext instanceContext; - - private boolean isSharedBudget() { - return instanceContext != null; - } - - private long getMaxBytesCanReserve(PlanNodeReservation reservation, long legacyLimit) { - return isSharedBudget() - ? (reservation.side == MemorySide.INPUT - ? instanceContext.getInputMaxBytes() - : instanceContext.getOutputMaxBytes()) - : legacyLimit; - } + // Enabled by local planning before drivers start; request limits are supplied by handles. + private boolean sharedBudget; private AtomicLong getUsedBytes(MemorySide side) { return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; @@ -217,7 +203,6 @@ public MemoryPool(String id, MemoryManager memoryManager, long maxBytesPerFragme .EXCEPTION_MAX_BYTES_PER_FI_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTES_RIGHT_BRACKET_DOT_MAXB_4D37C457, maxBytesPerFragmentInstance, this.memoryBlock.getTotalMemorySizeInBytes()); - this.maxBytesPerFragmentInstance = maxBytesPerFragmentInstance; } public String getId() { @@ -269,17 +254,15 @@ public void registerPlanNodeIdToQueryMemoryMap( } } - /** Register the FI-owned shared limits after local planning and before drivers start. */ - public void registerFragmentInstanceSharedBudget(FragmentInstanceContext instanceContext) { - String queryId = instanceContext.getId().getQueryId().getId(); - String fragmentInstanceId = instanceContext.getId().getFullId(); + /** Enable shared accounting after the complete FI has been checked by local planning. */ + public void registerFragmentInstanceSharedBudget(String queryId, String fragmentInstanceId) { synchronized (queryMemoryReservations) { FragmentMemoryReservation fragment = queryMemoryReservations .computeIfAbsent(queryId, x -> new ConcurrentHashMap<>()) .computeIfAbsent(fragmentInstanceId, x -> new FragmentMemoryReservation()); synchronized (fragment) { - fragment.instanceContext = instanceContext; + fragment.sharedBudget = true; } } } @@ -364,19 +347,10 @@ public MemoryReservationResult reserveWithPriority( Validate.notNull( planNodeId, DataNodeQueryMessages.EXCEPTION_PLANNODEID_CAN_NOT_BE_NULL_DOT_44027620); Validate.isTrue( - bytesToReserve > 0L && bytesToReserve <= maxBytesPerFragmentInstance, - DataNodeQueryMessages - .EXCEPTION_BYTESTORESERVE_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTESPERFI_RIGHT_BRACKET_DOT_M_0753BB69, - maxBytesPerFragmentInstance, - bytesToReserve); - FragmentMemoryReservation fragment = - getFragmentReservation(queryId, fragmentInstanceId, planNodeId); - synchronized (fragment) { - PlanNodeReservation reservation = - requireReservation( - fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); - maxBytesCanReserve = fragment.getMaxBytesCanReserve(reservation, maxBytesCanReserve); - } + bytesToReserve > 0L, + DataNodeQueryMessages.EXCEPTION_BYTESTORESERVE_SHOULD_BE_GREATER_THAN_ZERO_DOT_56D15DE0); + // A request must fit both the caller's FI/handle limit and the physical global pool. + maxBytesCanReserve = Math.min(maxBytesCanReserve, getMaxBytes()); if (bytesToReserve > maxBytesCanReserve) { LOGGER.warn( DataNodeQueryMessages @@ -430,11 +404,8 @@ public boolean tryReserveForTest( Validate.notNull( planNodeId, DataNodeQueryMessages.EXCEPTION_PLANNODEID_CAN_NOT_BE_NULL_DOT_44027620); Validate.isTrue( - bytesToReserve > 0L && bytesToReserve <= maxBytesPerFragmentInstance, - DataNodeQueryMessages - .EXCEPTION_BYTESTORESERVE_SHOULD_BE_IN_LEFT_PAREN_0_COMMA_MAXBYTESPERFI_RIGHT_BRACKET_DOT_M_0753BB69, - maxBytesPerFragmentInstance, - bytesToReserve); + bytesToReserve > 0L, + DataNodeQueryMessages.EXCEPTION_BYTESTORESERVE_SHOULD_BE_GREATER_THAN_ZERO_DOT_56D15DE0); return tryReserve(queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); } @@ -573,9 +544,8 @@ public boolean tryReserve( PlanNodeReservation reservation = requireReservation( fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); - maxBytesCanReserve = fragment.getMaxBytesCanReserve(reservation, maxBytesCanReserve); long usedBytes = - fragment.isSharedBudget() + fragment.sharedBudget ? fragment.getUsedBytes(reservation.side).get() : reservation.bytesReserved; // Do not publish speculative charges. Two competing requests must not reject one another diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index d7e89cbd6ff3..f5ecc27fe45a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -297,12 +297,15 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { instanceContext.initializeExchangeMemoryBudget( sharedBudget ? inputMaxBytes : 0, sharedBudget ? outputMaxBytes : 0); if (sharedBudget) { - memoryPool.registerFragmentInstanceSharedBudget(instanceContext); + memoryPool.registerFragmentInstanceSharedBudget( + instanceContext.getId().getQueryId().getId(), instanceContext.getId().getFullId()); } legacySinkLimits.forEach( (sink, legacyLimit) -> - sink.setMaxBytesCanReserve(sharedBudget ? outputMaxBytes : legacyLimit)); - long sourceLimit = sharedBudget ? inputMaxBytes : getMaxBytesOneHandleCanReserve(); + sink.setMaxBytesCanReserve( + sharedBudget ? instanceContext.getOutputMaxBytes() : legacyLimit)); + long sourceLimit = + sharedBudget ? instanceContext.getInputMaxBytes() : getMaxBytesOneHandleCanReserve(); exchangeOperatorList.forEach( exchangeOperator -> exchangeOperator.getSourceHandle().setMaxBytesCanReserve(sourceLimit)); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 1744216249f8..ed19981b400d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -19,13 +19,8 @@ package org.apache.iotdb.db.queryengine.execution.memory; -import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.commons.memory.MemoryManager; import org.apache.iotdb.db.exception.runtime.MemoryLeakException; -import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; -import org.apache.iotdb.db.queryengine.common.QueryId; -import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; -import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool.MemoryReservationResult; import com.google.common.util.concurrent.ListenableFuture; @@ -66,20 +61,6 @@ public void before() { QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, MemorySide.INPUT); } - private void registerSharedBudget( - MemoryPool target, String fragmentInstanceId, long inputMaxBytes, long outputMaxBytes) { - FragmentInstanceId id = Mockito.mock(FragmentInstanceId.class); - Mockito.when(id.getQueryId()).thenReturn(new QueryId(QUERY_ID)); - Mockito.when(id.getFullId()).thenReturn(fragmentInstanceId); - FragmentInstanceContext instanceContext = - FragmentInstanceContext.createFragmentInstanceContext( - id, - Mockito.mock(FragmentInstanceStateMachine.class), - Mockito.mock(MemoryReservationManager.class)); - instanceContext.initializeExchangeMemoryBudget(inputMaxBytes, outputMaxBytes); - target.registerFragmentInstanceSharedBudget(instanceContext); - } - @Test public void testDeregisterDoesNotRemoveReregisteredQueryMemory() throws Exception { Field reservationsField = MemoryPool.class.getDeclaredField("queryMemoryReservations"); @@ -548,7 +529,7 @@ public void testReserveWithPriorityShowQueriesWithinAvailable() { @Test public void testSharedInputBudgetAndIndependentOutput() { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -595,12 +576,33 @@ public void testSideRegistrationCannotChange() { } @Test - public void testSharedRequestCannotExceedInitializedInputBudget() { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + public void testRequestsUseCurrentSuppliedLimitInsteadOfPoolCreationBudget() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + Assert.assertTrue( + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 768, 768, false) + .isReserveSuccess()); + Assert.assertEquals(768, pool.getReservedBytes()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 768); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testRequestLargerThanGlobalCapacityDoesNotWait() { try { - pool.reserveWithPriority( - QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, Long.MAX_VALUE, false); - Assert.fail("request exceeds the initialized input budget"); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 1025, 2048, false); + Assert.fail("a request larger than the global pool cannot become ready"); + } catch (IllegalArgumentException expected) { + Assert.assertEquals(0, pool.getReservedBytes()); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + } + } + + @Test + public void testSharedRequestCannotExceedSuppliedInputBudget() { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + try { + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 257, 256, false); + Assert.fail("request exceeds the supplied input budget"); } catch (IllegalArgumentException expected) { Assert.assertEquals(0, pool.getReservedBytes()); Assert.assertEquals(0, pool.getQueryMemoryReservedBytes(QUERY_ID)); @@ -609,27 +611,25 @@ public void testSharedRequestCannotExceedInitializedInputBudget() { } @Test - public void testHandleLimitsCannotChangeInitializedSharedBudgets() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + public void testSharedReservationsAndWakeupsUseSuppliedSideLimits() throws Exception { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "output", MemorySide.OUTPUT); Assert.assertTrue( - pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 1, false) + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200, 256, false) .isReserveSuccess()); MemoryReservationResult pending = - pool.reserveWithPriority( - QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, Long.MAX_VALUE, false); + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256, false); Assert.assertFalse(pending.getFuture().isDone()); Assert.assertTrue( - pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 1, false) + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256, 256, false) .isReserveSuccess()); - // Waking the input request must still respect the FI's saved limit. + // An OUTPUT release cannot grant an INPUT request whose shared side remains full. pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, "output", 256); Assert.assertFalse(pending.getFuture().isDone()); - Assert.assertFalse( - pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, Long.MAX_VALUE)); + Assert.assertFalse(pool.tryReserveForTest(QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", 100, 256)); pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 200); pending.getFuture().get(10, TimeUnit.SECONDS); Assert.assertEquals( @@ -645,21 +645,17 @@ public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { new MemoryPool("single-direction", new MemoryManager(1024), 513); for (MemorySide side : MemorySide.values()) { String fi = side.name(); - registerSharedBudget( - singleDirectionPool, - fi, - side == MemorySide.INPUT ? 513 : 0, - side == MemorySide.OUTPUT ? 513 : 0); + singleDirectionPool.registerFragmentInstanceSharedBudget(QUERY_ID, fi); singleDirectionPool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, "active", side); singleDirectionPool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, fi, "unused", side == MemorySide.INPUT ? MemorySide.OUTPUT : MemorySide.INPUT); Assert.assertTrue( singleDirectionPool - .reserveWithPriority(QUERY_ID, fi, "active", 513, 1, false) + .reserveWithPriority(QUERY_ID, fi, "active", 513, 513, false) .isReserveSuccess()); try { - singleDirectionPool.reserveWithPriority(QUERY_ID, fi, "unused", 1, Long.MAX_VALUE, false); - Assert.fail("an unused direction must not acquire a budget from a handle"); + singleDirectionPool.reserveWithPriority(QUERY_ID, fi, "unused", 1, 0, false); + Assert.fail("an unused direction supplies a zero limit"); } catch (IllegalArgumentException expected) { Assert.assertEquals(0, singleDirectionPool.getMemoryReservationSize()); Assert.assertEquals(513, singleDirectionPool.getReservedBytes()); @@ -671,7 +667,7 @@ public void testSingleDirectionKeepsFullBudgetAndSingleBlockLimit() { @Test public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 128, 128); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); ExecutorService executor = Executors.newFixedThreadPool(16); CountDownLatch start = new CountDownLatch(1); CountDownLatch attempted = new CountDownLatch(16); @@ -719,7 +715,7 @@ public void testConcurrentPlanNodesShareOneSideCounter() throws Exception { @Test public void testReleaseBeforePendingReservationIsPublished() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool = Mockito.spy(pool); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "input2", MemorySide.INPUT); @@ -772,7 +768,7 @@ public void testReleaseBeforePendingReservationIsPublished() throws Exception { @Test public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap( QUERY_ID, FRAGMENT_INSTANCE_ID, "cancelled", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap( @@ -810,7 +806,7 @@ public void testCancelledInputDoesNotConsumeReleasedSharedBudget() throws Except @Test public void testConcurrentPendingReservationsCompleteAndReleaseAllAccounting() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 128, 128); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); int count = 16; ExecutorService executor = Executors.newFixedThreadPool(count); CountDownLatch attempted = new CountDownLatch(count); @@ -871,7 +867,7 @@ public void testDefaultFragmentKeepsIndependentPlanNodeQuotas() { @Test public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "a", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "legacy", "b", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, "b", MemorySide.INPUT); @@ -895,7 +891,7 @@ public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { @Test public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { - registerSharedBudget(pool, FRAGMENT_INSTANCE_ID, 256, 256); + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); for (int i = 0; i < 3; i++) { String fi = "other" + i; pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 21dafba0432e..ebcc275717f6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -186,10 +186,12 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit @Test public void testPureCollectUsesSharedSideBudgets() { ISourceHandle source = addSource(context); + ISourceHandle secondSource = addSource(context); ISinkHandle sink = mock(ISinkHandle.class); context.setISink(sink); context.setExchangeMemoryLimits(node(PlanNodeType.IDENTITY_SINK, collect()), pool); verify(source).setMaxBytesCanReserve(512); + verify(secondSource).setMaxBytesCanReserve(512); verify(sink).setMaxBytesCanReserve(512); assertAdmissionIsShared(true, MemorySide.INPUT, 512); assertAdmissionIsShared(true, MemorySide.OUTPUT, 512); diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java index 82fd7cc80e22..be7cf49239a6 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/ClientManagerTest.java @@ -55,8 +55,7 @@ public class ClientManagerTest { - private final TEndPoint endPoint = - new TEndPoint("localhost", Integer.getInteger("iotdb.test.client-manager.port", 10730)); + private final TEndPoint endPoint = new TEndPoint("localhost", 10730); private static final int CONNECTION_TIMEOUT = 5_000; From 3d7ebb95425dcb9d99686ce98637bea3850fb2d1 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 18:44:40 +0800 Subject: [PATCH 11/15] Avoid per-block Collect input scans with a ready queue --- .../operator/process/AnyChildBlocked.java | 180 +++++++++++------- .../operator/process/CollectOperatorTest.java | 163 +++++++++++++++- 2 files changed, 270 insertions(+), 73 deletions(-) diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java index c0397bf1ff5a..c1ef2edb1b61 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java @@ -50,10 +50,13 @@ final class AnyChildBlocked implements Accountable { private final List children; private final Object lock = new Object(); private final ChildFutureListener[] listeners; + private ChildFutureListener readyHead; + private ChildFutureListener readyTail; private WaitRound currentRound; private Throwable firstFailure; private boolean closed; - private int nextChild; + private boolean initialized; + private int childToObserve = -1; AnyChildBlocked(List children) { this.children = children; @@ -62,65 +65,100 @@ final class AnyChildBlocked implements Accountable { ListenableFuture getAnyChild() { WaitRound round; + WaitRound selection; + boolean initialize; + int invalidatedChild; synchronized (lock) { if (currentRound != null) { return currentRound.selected; } round = new WaitRound(); currentRound = round; - for (ChildFutureListener listener : listeners) { - if (listener != null) { - listener.target = round; - } + initialize = !initialized; + initialized = true; + invalidatedChild = childToObserve; + childToObserve = -1; + selection = selectReadyChild(); + } + completeSelection(selection); + + if (initialize) { + // Activate every direct input once, even if an earlier input is already ready. Pending + // listeners enqueue their completion without being retargeted or polled on every block. + for (int childId = 0; childId < children.size(); childId++) { + observeChild(childId); } + } else if (invalidatedChild >= 0) { + // Only advancing an input invalidates its readiness. Requeue it behind inputs that are + // already ready so that a continuously ready input cannot starve its siblings. + observeChild(invalidatedChild); } + return round.selected; + } - // Activate every live direct child, even if an earlier one has already become ready. Only - // invalidated records call child.isBlocked(); a long-pending future keeps its one listener. - for (int offset = 0; offset < children.size(); offset++) { - int childId = (nextChild + offset) % children.size(); - Operator child; - ChildFutureListener listener; - synchronized (lock) { - if (closed) { - break; - } - child = children.get(childId); - listener = listeners[childId]; + private void observeChild(int childId) { + Operator child; + synchronized (lock) { + if (closed) { + return; } - if (child == null) { - continue; + child = children.get(childId); + } + if (child == null) { + return; + } + ListenableFuture observedFuture = child.isBlocked(); + ChildFutureListener listener = new ChildFutureListener(childId, observedFuture); + synchronized (lock) { + if (closed) { + return; } - if (listener == null) { - ListenableFuture observedFuture = child.isBlocked(); - synchronized (lock) { - if (closed) { - break; - } - listener = new ChildFutureListener(childId, observedFuture, round); - listeners[childId] = listener; + listeners[childId] = listener; + } + // Publish the record before registration, including when completion runs inline. Callbacks + // only update readiness; operator methods remain confined to the driver thread. + observedFuture.addListener(listener, directExecutor()); + } + + /** Called under lock; claiming a selection and removing its ready record are atomic. */ + private WaitRound selectReadyChild() { + if (currentRound == null || currentRound.claimed) { + return null; + } + if (closed || children.isEmpty()) { + currentRound.childId = -1; + } else if (firstFailure != null) { + currentRound.failure = firstFailure; + } else { + ChildFutureListener listener; + do { + listener = readyHead; + if (listener == null) { + return null; } - // The target is published before registration. addListener also handles a future that - // finishes between isBlocked and registration, including inline completion. - observedFuture.addListener(listener, directExecutor()); - } else if (listener.observedFuture.isDone()) { - // A completed record may have notified a previous round. Recheck it for this round. - listener.run(); - } + readyHead = listener.nextReady; + listener.nextReady = null; + if (readyHead == null) { + readyTail = null; + } + // Retired records are discarded once, rather than scanning closed input slots per block. + } while (listeners[listener.childId] != listener); + currentRound.childId = listener.childId; } + currentRound.claimed = true; + return currentRound; + } - Throwable failure; - boolean isClosed; - synchronized (lock) { - failure = firstFailure; - isClosed = closed; + private static void completeSelection(WaitRound selection) { + if (selection == null) { + return; } - if (isClosed || children.isEmpty()) { - round.selected.set(-1); - } else if (failure != null) { - round.selected.setException(failure); + // Completing the future can invoke downstream listeners synchronously. Do not hold lock. + if (selection.failure == null) { + selection.selected.set(selection.childId); + } else { + selection.selected.setException(selection.failure); } - return round.selected; } void checkFailure() throws Exception { @@ -142,11 +180,8 @@ void checkFailure() throws Exception { void invalidateChild(int childId) { synchronized (lock) { listeners[childId] = null; - if (currentRound != null) { - currentRound.valid = false; - currentRound = null; - } - nextChild = (childId + 1) % children.size(); + currentRound = null; + childToObserve = childId; } } @@ -160,14 +195,13 @@ void close() { closed = true; round = currentRound; currentRound = null; - if (round != null) { - round.valid = false; - } - for (ChildFutureListener listener : listeners) { - if (listener != null) { - listener.target = null; - } + childToObserve = -1; + while (readyHead != null) { + ChildFutureListener next = readyHead.nextReady; + readyHead.nextReady = null; + readyHead = next; } + readyTail = null; Arrays.fill(listeners, null); } if (round != null) { @@ -193,18 +227,20 @@ public long ramBytesUsed() { private static final class WaitRound { private final SettableFuture selected = SettableFuture.create(); - private boolean valid = true; + private boolean claimed; + private int childId; + private Throwable failure; } private final class ChildFutureListener implements Runnable { private final int childId; private final ListenableFuture observedFuture; - private WaitRound target; + private ChildFutureListener nextReady; + private boolean completed; - private ChildFutureListener(int childId, ListenableFuture observedFuture, WaitRound target) { + private ChildFutureListener(int childId, ListenableFuture observedFuture) { this.childId = childId; this.observedFuture = observedFuture; - this.target = target; } @Override @@ -217,26 +253,26 @@ public void run() { } catch (CancellationException e) { failure = e; } - WaitRound round; + WaitRound selection; synchronized (lock) { - if (closed || listeners[childId] != this) { + if (closed || listeners[childId] != this || completed) { return; } + completed = true; if (failure != null && firstFailure == null) { firstFailure = failure; } - round = target; - if (round == null || !round.valid || round != currentRound) { - return; + if (failure == null) { + if (readyTail == null) { + readyHead = this; + } else { + readyTail.nextReady = this; + } + readyTail = this; } - failure = firstFailure; - } - // Completing a selection can invoke downstream listeners synchronously. Do not hold lock. - if (failure == null) { - round.selected.set(childId); - } else { - round.selected.setException(failure); + selection = selectReadyChild(); } + completeSelection(selection); } } } diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java index b446f9c999b9..6de8128fe70a 100644 --- a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -33,6 +33,7 @@ import org.junit.Test; import java.io.IOException; +import java.util.AbstractList; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -47,6 +48,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertEquals; @@ -120,7 +122,98 @@ public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exceptio } @Test - public void testPendingListenerRetargetedAndStaleCompletionIgnored() throws Exception { + public void testAllReadyChildrenHaveLinearSelectionCost() throws Exception { + int childCount = 1000; + int rounds = 5; + List children = new ArrayList<>(); + for (int i = 0; i < childCount; i++) { + Child child = new Child(block(i)); + child.ready(); + children.add(child); + } + CountingList countedChildren = new CountingList(children); + CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + try { + for (int round = 0; round < rounds; round++) { + assertTrue(collect.isBlocked().isDone()); + // A completed future may notify more than once. It must occupy only one ready slot. + for (Child child : children) { + child.wait.listener.run(); + child.wait.listener.run(); + } + Set consumed = new HashSet<>(); + for (int i = 0; i < childCount; i++) { + assertTrue(collect.isBlocked().isDone()); + TsBlock result = collect.next(); + assertTrue( + "every ready child must be selected once per round", + consumed.add(result.getColumn(0).getInt(0))); + } + assertEquals(childCount, consumed.size()); + } + for (Child child : children) { + assertEquals(rounds, child.advances); + } + long probes = countedChildren.getCalls + countFutureProbes(children); + long linearLimit = 20L * (childCount + childCount * rounds); + assertTrue( + "selection must use O(children + blocks) probes, observed " + probes, + probes <= linearLimit); + } finally { + collect.close(); + } + } + + @Test + public void testClosedChildrenDoNotIncreaseRemainingChildSelectionCost() throws Exception { + int childCount = 1000; + int remainingBlocks = 3000; + List children = new ArrayList<>(); + for (int i = 0; i < childCount; i++) { + Child child = new Child(block(i)); + child.hasNext = i == childCount - 1; + child.finished = !child.hasNext; + child.ready(); + children.add(child); + } + CountingList countedChildren = new CountingList(children); + CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + try { + // Retire the first 999 children and consume one block from the remaining child. + for (int i = 0; i < childCount; i++) { + assertTrue(collect.isBlocked().isDone()); + collect.next(); + } + for (int i = 0; i < childCount - 1; i++) { + assertEquals(1, children.get(i).closes); + } + Child remaining = children.get(childCount - 1); + assertEquals(1, remaining.advances); + long initialProbes = countedChildren.getCalls + countFutureProbes(children); + for (int i = 0; i < remainingBlocks; i++) { + assertTrue(collect.isBlocked().isDone()); + assertSame(remaining.output, collect.next()); + } + assertEquals(remainingBlocks + 1, remaining.advances); + long probes = countedChildren.getCalls + countFutureProbes(children) - initialProbes; + assertTrue( + "retired children must not be visited for each remaining block, observed " + probes, + probes <= 20L * remainingBlocks); + } finally { + collect.close(); + } + } + + private static long countFutureProbes(List children) { + long probes = 0; + for (Child child : children) { + probes += child.wait.probes.get(); + } + return probes; + } + + @Test + public void testPendingListenerReusedAndStaleCompletionIgnored() throws Exception { Child a = new Child(block(1)); Child b = new Child(block(2)); CollectOperator collect = collect(a, b); @@ -141,6 +234,29 @@ public void testPendingListenerRetargetedAndStaleCompletionIgnored() throws Exce collect.close(); } + @Test + public void testCompletionBetweenSelectionsIsNotLost() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = collect(a, b); + try { + assertTrue(collect.isBlocked().isDone()); + assertSame(b.output, collect.next()); + b.reset(); + // No selection is waiting yet: retain this completion until the next driver observation. + a.ready(); + assertTrue(collect.isBlocked().isDone()); + assertSame(a.output, collect.next()); + a.reset(); + assertFalse(collect.isBlocked().isDone()); + assertEquals(1, a.advances); + assertEquals(1, b.advances); + } finally { + collect.close(); + } + } + @Test public void testNullEmptyAndUnfinishedChildRefreshReadiness() throws Exception { Child child = new Child(null); @@ -401,10 +517,36 @@ public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Excep } } + private static final class CountingList extends AbstractList { + private final List delegate; + private long getCalls; + + private CountingList(List children) { + delegate = new ArrayList<>(children); + } + + @Override + public Operator get(int index) { + getCalls++; + return delegate.get(index); + } + + @Override + public Operator set(int index, Operator child) { + return delegate.set(index, child); + } + + @Override + public int size() { + return delegate.size(); + } + } + private static final class CountingFuture extends ForwardingListenableFuture.SimpleForwardingListenableFuture { private final SettableFuture delegate; private final AtomicInteger registrations = new AtomicInteger(); + private final AtomicInteger probes = new AtomicInteger(); private Runnable beforeRegistration; private Runnable listener; @@ -417,6 +559,25 @@ private CountingFuture(SettableFuture delegate) { this.delegate = delegate; } + @Override + public boolean isDone() { + probes.incrementAndGet(); + return super.isDone(); + } + + @Override + public Void get() throws InterruptedException, ExecutionException { + probes.incrementAndGet(); + return super.get(); + } + + @Override + public Void get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + probes.incrementAndGet(); + return super.get(timeout, unit); + } + @Override public void addListener(Runnable listener, Executor executor) { this.listener = listener; From b50d928c85474d6b3432a62cb344193d1e7616d2 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 15 Oct 2026 20:40:53 +0800 Subject: [PATCH 12/15] Limit shared exchange budgets and ready Collect to table fragments --- .../it/query/recent/IoTDBCollectFanInIT.java | 21 +++- .../operator/process/CollectOperator.java | 61 ++++++++++- .../operator/process/CollectOperatorTest.java | 51 ++++++++- .../iotdb/db/i18n/DataNodeQueryMessages.java | 1 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 1 + .../DataNodeTableOperatorGenerator.java | 17 +++ .../planner/LocalExecutionPlanContext.java | 44 +++++++- .../plan/planner/LocalExecutionPlanner.java | 38 ++++--- .../plan/planner/OperatorTreeGenerator.java | 7 +- .../SharedExchangeMemoryEligibility.java | 19 ++-- .../TableDistributedPlanGenerator.java | 25 +---- .../LocalExecutionPlanContextBudgetTest.java | 103 +++++++++++++++++- .../planner/CollectInputPropertyTest.java | 19 ---- 13 files changed, 322 insertions(+), 85 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java index 943a2150d7f6..356908b5affc 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -34,7 +34,9 @@ import java.sql.ResultSet; import java.sql.Statement; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.assertUnorderedQueryResult; import static org.junit.Assert.assertEquals; @@ -109,7 +111,7 @@ public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Excepti } @Test - public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception { + public void testLegacyTreeAndSharedTreeToTablePaths() throws Exception { try (Connection connection = EnvFactory.getEnv().getConnection(); Statement statement = connection.createStatement()) { statement.execute("CREATE DATABASE root.collect_tree"); @@ -124,10 +126,25 @@ public void testSharedExchangePathsForTreeAndTreeToTable() throws Exception { try (ResultSet result = statement.executeQuery("SELECT value FROM root.collect_tree.*.* ALIGN BY DEVICE")) { int rows = 0; + Set seenDevices = new HashSet<>(); + String currentDevice = null; + int deviceRows = 0; while (result.next()) { + String nextDevice = result.getString("Device"); + if (!nextDevice.equals(currentDevice)) { + if (currentDevice != null) { + assertEquals(ROWS, deviceRows); + } + assertTrue("Tree results must retain contiguous devices", seenDevices.add(nextDevice)); + currentDevice = nextDevice; + deviceRows = 0; + } + assertEquals(++deviceRows, result.getLong("Time")); assertEquals(result.getLong("Time"), result.getLong("value")); rows++; } + assertEquals(ROWS, deviceRows); + assertEquals(DEVICES, seenDevices.size()); assertEquals(DEVICES * ROWS, rows); } } @@ -300,7 +317,7 @@ public void testOrderByAfterGlobalWindowAndConstantFill() throws Exception { } @Test - public void testTopKAfterGlobalWindowConsumesUnorderedInput() throws Exception { + public void testExcludedTopKPlanRetainsQueryResults() throws Exception { try (Connection connection = EnvFactory.getEnv().getTableConnection(); Statement statement = connection.createStatement()) { statement.execute("USE " + DATABASE); diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java index bd7dec028b13..c6ee9577efe6 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java @@ -39,7 +39,9 @@ public class CollectOperator implements ProcessOperator { private final CommonOperatorContext operatorContext; protected final List children; - private final AnyChildBlocked anyChildBlocked; + private AnyChildBlocked anyChildBlocked; + private boolean inited; + protected int currentIndex; private ListenableFuture selectionFuture; private ListenableFuture exposedFuture; private int remainingChildren; @@ -48,17 +50,36 @@ public class CollectOperator implements ProcessOperator { public CollectOperator(CommonOperatorContext operatorContext, List children) { this.operatorContext = operatorContext; this.children = children; - this.anyChildBlocked = new AnyChildBlocked(children); this.remainingChildren = children.size(); } + /** + * Enable only after the complete table FI has passed eligibility checks, before drivers start. + */ + public void enableReadyInputConsumption() { + if (anyChildBlocked == null) { + anyChildBlocked = new AnyChildBlocked(children); + } + } + @Override public boolean hasNext() { - return !closed && remainingChildren > 0; + return anyChildBlocked == null + ? currentIndex < children.size() + : !closed && remainingChildren > 0; } @Override public TsBlock next() throws Exception { + if (anyChildBlocked == null) { + if (children.get(currentIndex).hasNextWithTimer()) { + TsBlock block = children.get(currentIndex).nextWithTimer(); + return block == null ? null : transformBlock(currentIndex, block); + } + closeCurrentChild(currentIndex); + currentIndex++; + return null; + } anyChildBlocked.checkFailure(); if (!hasNext()) { return null; @@ -98,6 +119,11 @@ protected TsBlock transformBlock(int childId, TsBlock block) { } protected void closeCurrentChild(int index) throws Exception { + if (anyChildBlocked == null) { + children.get(index).close(); + children.set(index, null); + return; + } Operator child = children.get(index); if (child == null) { return; @@ -111,6 +137,15 @@ protected void closeCurrentChild(int index) throws Exception { @Override public ListenableFuture isBlocked() { + if (anyChildBlocked == null) { + if (!inited) { + inited = true; + for (Operator child : children) { + child.isBlocked(); + } + } + return currentIndex >= children.size() ? NOT_BLOCKED : children.get(currentIndex).isBlocked(); + } if (!hasNext()) { return NOT_BLOCKED; } @@ -135,6 +170,14 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { + if (anyChildBlocked == null) { + for (int i = currentIndex; i < children.size(); i++) { + if (children.get(i) != null) { + closeCurrentChild(i); + } + } + return; + } if (closed) { return; } @@ -161,6 +204,13 @@ public void close() throws Exception { @Override public long calculateMaxPeekMemory() { + if (anyChildBlocked == null) { + long maxPeekMemory = 0; + for (Operator child : children) { + maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + } + return maxPeekMemory; + } long retained = calculateRetainedSizeAfterCallingNext(); long peak = retained; for (Operator child : children) { @@ -189,6 +239,9 @@ public long calculateMaxReturnSize() { @Override public long calculateRetainedSizeAfterCallingNext() { + if (anyChildBlocked == null) { + return 0; + } long retained = 0; for (Operator child : children) { if (child != null) { @@ -206,7 +259,7 @@ public List getChildren() { @Override public long ramBytesUsed() { return INSTANCE_SIZE - + anyChildBlocked.ramBytesUsed() + + (anyChildBlocked == null ? 0 : anyChildBlocked.ramBytesUsed()) + children.stream() .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject) .sum() diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java index 6de8128fe70a..6ab0b3a35f92 100644 --- a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -71,7 +71,10 @@ private static TsBlock block(int... values) { } private static CollectOperator collect(Operator... children) { - return new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children))); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(children))); + collect.enableReadyInputConsumption(); + return collect; } @Test @@ -100,6 +103,7 @@ public void testThousandChildrenAndOneListenerPerPendingPeriod() throws Exceptio children.add(new Child(block(i))); } CollectOperator collect = new CollectOperator(new TestContext(), children); + collect.enableReadyInputConsumption(); Child ready = (Child) children.get(count - 1); for (int round = 0; round < 5; round++) { ListenableFuture wait = collect.isBlocked(); @@ -133,6 +137,7 @@ public void testAllReadyChildrenHaveLinearSelectionCost() throws Exception { } CountingList countedChildren = new CountingList(children); CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + collect.enableReadyInputConsumption(); try { for (int round = 0; round < rounds; round++) { assertTrue(collect.isBlocked().isDone()); @@ -178,6 +183,7 @@ public void testClosedChildrenDoNotIncreaseRemainingChildSelectionCost() throws } CountingList countedChildren = new CountingList(children); CollectOperator collect = new CollectOperator(new TestContext(), countedChildren); + collect.enableReadyInputConsumption(); try { // Retire the first 999 children and consume one block from the remaining child. for (int i = 0; i < childCount; i++) { @@ -378,6 +384,47 @@ public void testCloseWakesWaitAndClosesEveryChildDespiteException() throws Excep assertEquals(1, b.closes); } + @Test + public void testDefaultCollectRetainsSequentialConsumption() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertFalse(collect.isBlocked().isDone()); + a.ready(); + assertSame(a.output, collect.next()); + a.hasNext = false; + a.finished = true; + assertNull(collect.next()); + assertSame(b.output, collect.next()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + collect.close(); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testDefaultMappingCollectRetainsSequentialConsumption() throws Exception { + Child a = new Child(block(10, 11)); + Child b = new Child(block(20, 21)); + b.ready(); + List> mappings = new ArrayList<>(); + mappings.add(Collections.singletonList(0)); + mappings.add(Collections.singletonList(1)); + MappingCollectOperator collect = + new MappingCollectOperator( + new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings); + assertFalse(collect.isBlocked().isDone()); + a.ready(); + assertEquals(10, collect.next().getColumn(0).getInt(0)); + a.hasNext = false; + a.finished = true; + assertNull(collect.next()); + assertEquals(21, collect.next().getColumn(0).getInt(0)); + collect.close(); + } + @Test public void testMappingUsesActualChildIdentity() throws Exception { Child a = new Child(block(10, 11)); @@ -389,6 +436,7 @@ public void testMappingUsesActualChildIdentity() throws Exception { MappingCollectOperator collect = new MappingCollectOperator( new TestContext(), new ArrayList<>(Arrays.asList(a, b)), mappings); + collect.enableReadyInputConsumption(); assertTrue(collect.isBlocked().isDone()); assertEquals(21, collect.next().getColumn(0).getInt(0)); a.ready(); @@ -473,6 +521,7 @@ public void testConcurrentCompletionsConsumeEveryChildExactlyOnce() throws Excep List children = new ArrayList<>(); for (int i = 0; i < count; i++) children.add(new Child(block(i))); CollectOperator collect = new CollectOperator(new TestContext(), new ArrayList<>(children)); + collect.enableReadyInputConsumption(); ListenableFuture pending = collect.isBlocked(); assertFalse(pending.isDone()); ExecutorService executor = Executors.newFixedThreadPool(count); 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 be8c3ade6477..965ae20034cb 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 @@ -3943,4 +3943,5 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "Plan node %s is already registered on memory side %s, not %s"; public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "Memory reservation is not registered for query %s, fragment %s, plan node %s"; + public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "Shared exchange budget for FI {}: INPUT={}, OUTPUT={}"; } 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 13fa4222a40d..290504411081 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 @@ -4699,4 +4699,5 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_PLAN_NODE_ARG_IS_ALREADY_REGISTERED_ON_MEMORY_SIDE_ARG_NOT_ARG_4C1BB625 = "PlanNode %s 已注册到内存侧 %s，不能注册到 %s"; public static final String EXCEPTION_MEMORY_RESERVATION_IS_NOT_REGISTERED_FOR_QUERY_ARG_FRAGMENT_ARG_PLAN_NODE_ARG_67A4B81F = "查询 %s、Fragment %s、PlanNode %s 尚未注册内存账目"; + public static final String LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109 = "FI {} 的共享 exchange 预算：INPUT={}，OUTPUT={}"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java index 2fae9aa53708..4b4ca52b8558 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/DataNodeTableOperatorGenerator.java @@ -22,6 +22,7 @@ import org.apache.iotdb.calc.execution.filter.TopKRuntimeFilter; import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.execution.operator.process.FilterAndProjectOperator; import org.apache.iotdb.calc.execution.operator.process.LimitOperator; import org.apache.iotdb.calc.execution.operator.process.OffsetOperator; @@ -47,7 +48,9 @@ import org.apache.iotdb.commons.queryengine.plan.relational.metadata.QualifiedObjectName; import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.AggregationNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.LongLiteral; @@ -301,6 +304,20 @@ protected OperatorContext addOperatorContext( .addOperatorContext(context.getNextOperatorId(), planNodeId, operatorType); } + @Override + public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context) { + CollectOperator collect = (CollectOperator) super.visitCollect(node, context); + context.addCollectOperator(collect); + return collect; + } + + @Override + public Operator visitUnion(UnionNode node, LocalExecutionPlanContext context) { + CollectOperator collect = (CollectOperator) super.visitUnion(node, context); + context.addCollectOperator(collect); + return collect; + } + @Override public Operator visitCteScan(CteScanNode node, LocalExecutionPlanContext context) { OperatorContext operatorContext = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index f5ecc27fe45a..9e5cd2e0ea00 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.plan.planner.ITableOperatorGeneratorContext; import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.calc.plan.planner.memory.PipelineMemoryEstimator; @@ -51,6 +52,8 @@ import org.apache.tsfile.read.filter.basic.Filter; import org.apache.tsfile.utils.Binary; import org.apache.tsfile.utils.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import javax.annotation.Nullable; @@ -73,6 +76,8 @@ // Attention: We should use thread-safe data structure for members that are shared by all pipelines public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext { + private static final Logger LOGGER = LoggerFactory.getLogger(LocalExecutionPlanContext.class); + // Save operators in this pipeline, a new one will be created when creating another pipeline private final DriverContext driverContext; private final AtomicInteger nextOperatorId; @@ -87,6 +92,7 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext private int exchangeSumNum = 0; private final long fragmentExchangeMaxBytes; private Map legacySinkLimits = new IdentityHashMap<>(); + private List collectOperators = new ArrayList<>(); private List cachedDataTypes; @@ -129,6 +135,7 @@ public LocalExecutionPlanContext( public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) { this.fragmentExchangeMaxBytes = parentContext.fragmentExchangeMaxBytes; this.legacySinkLimits = parentContext.legacySinkLimits; + this.collectOperators = parentContext.collectOperators; this.nextOperatorId = parentContext.nextOperatorId; this.typeProvider = parentContext.typeProvider; this.allSensorsMap = parentContext.allSensorsMap; @@ -257,6 +264,7 @@ public void setExchangeSumNum(int exchangeSumNum) { this.exchangeSumNum = exchangeSumNum; } + /** Original per-handle quota for tree planning and independent table FI accounting. */ public long getMaxBytesOneHandleCanReserve() { return exchangeSumNum == 0 ? fragmentExchangeMaxBytes @@ -271,6 +279,23 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { this.exchangeOperatorList.add(exchangeOperator); } + public void addCollectOperator(CollectOperator collectOperator) { + collectOperators.add(collectOperator); + } + + public void setMaxBytesOneHandleCanReserve() { + long maxBytesOneHandleCanReserve = getMaxBytesOneHandleCanReserve(); + LOGGER.debug( + DataNodeQueryMessages + .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG, + maxBytesOneHandleCanReserve, + exchangeSumNum); + exchangeOperatorList.forEach( + exchangeOperator -> + exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve)); + } + + /** Called after operator generation, before root memory estimation and factory registration. */ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { // Only remote sources charge this FI's INPUT. Local inputs are owned by the producing FI; // internal pipeline queues cannot enable sharing because of the pipeline check below. @@ -288,9 +313,10 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { boolean sharedBudget = (inputMaxBytes > 0 || outputMaxBytes > 0) && SharedExchangeMemoryEligibility.isEligible(plan) - // Internal pipeline queues and their downstream sinks must not compete for OUTPUT. - && pipelineDriverFactories.size() == 1 - && pipelineDriverFactories.get(0).getDependencyPipelineIndex() == -1 + // Only the root pipeline may exist. Child factories are registered during generation; + // the root factory is registered after its memory has been estimated below. + && pipelineDriverFactories.isEmpty() + && driverContext.getDependencyDriverIndex() == -1 && (inputExchangeCount == 0 || inputMaxBytes >= maxBlockBytes) && (outputSinkCount == 0 || outputMaxBytes >= maxBlockBytes); FragmentInstanceContext instanceContext = driverContext.getFragmentInstanceContext(); @@ -299,6 +325,18 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { if (sharedBudget) { memoryPool.registerFragmentInstanceSharedBudget( instanceContext.getId().getQueryId().getId(), instanceContext.getId().getFullId()); + collectOperators.forEach(CollectOperator::enableReadyInputConsumption); + LOGGER.debug( + DataNodeQueryMessages.LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109, + instanceContext.getId(), + instanceContext.getInputMaxBytes(), + instanceContext.getOutputMaxBytes()); + } else { + LOGGER.debug( + DataNodeQueryMessages + .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG, + getMaxBytesOneHandleCanReserve(), + exchangeSumNum); } legacySinkLimits.forEach( (sink, legacyLimit) -> diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index 8e020cf73645..c600b4eb9016 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -116,6 +116,9 @@ public List plan( Operator root = generateOperator(instanceContext, context, plan); + // Select Collect consumption before estimating its retained input memory. + setExchangeMemoryLimits(plan, context); + PipelineMemoryEstimator memoryEstimator = context.constructPipelineMemoryEstimator(root, null, plan, -1); // set the map to null for gc @@ -133,14 +136,6 @@ public List plan( context.getTimePartitions().ifPresent(instanceContext::setTimePartitions); - // Choose the accounting policy only after all FI operators and pipelines are known. - context.setExchangeMemoryLimits( - plan, - MPPDataExchangeService.getInstance() - .getMPPDataExchangeManager() - .getLocalMemoryManager() - .getQueryPool()); - return context.getPipelineDriverFactories(); } @@ -161,6 +156,9 @@ public List plan( Operator root = generateOperator(instanceContext, context, plan); + // Select Collect consumption before estimating its retained input memory. + setExchangeMemoryLimits(plan, context); + PipelineMemoryEstimator memoryEstimator = context.constructPipelineMemoryEstimator(root, null, plan, -1); // set the map to null for gc @@ -171,14 +169,6 @@ public List plan( context.addPipelineDriverFactory(root, context.getDriverContext(), 0); - // Choose the accounting policy only after all FI operators and pipelines are known. - context.setExchangeMemoryLimits( - plan, - MPPDataExchangeService.getInstance() - .getMPPDataExchangeManager() - .getLocalMemoryManager() - .getQueryPool()); - return context.getPipelineDriverFactories(); } @@ -207,6 +197,22 @@ private Operator generateOperator( return root; } + private void setExchangeMemoryLimits(PlanNode plan, LocalExecutionPlanContext context) { + FragmentInstanceContext instanceContext = context.getInstanceContext(); + if (instanceContext.getSessionInfo() == null + || instanceContext.getSessionInfo().getSqlDialect() == TREE) { + context.setMaxBytesOneHandleCanReserve(); + return; + } + // Choose the table FI policy only after all operators and pipelines are known. + context.setExchangeMemoryLimits( + plan, + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool()); + } + private long checkMemory( final PipelineMemoryEstimator memoryEstimator, FragmentInstanceContext instanceContext) throws MemoryNotEnoughException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java index a46877ac1d5f..b6a57db5445d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java @@ -2577,6 +2577,7 @@ public Operator visitExchange(ExchangeNode node, LocalExecutionPlanContext conte if (!isSameNode) { context.addExchangeSumNum(1); } + sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); context.addExchangeOperator(exchangeOperator); @@ -2615,7 +2616,8 @@ public Operator visitIdentitySink(IdentitySinkNode node, LocalExecutionPlanConte node.getPlanNodeId().getId(), context.getInstanceContext()); List children = dealWithConsumeChildrenOneByOneNode(node, context); - context.setISink(sinkHandle); + sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + context.getDriverContext().setSink(sinkHandle); return new IdentitySinkOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); } @@ -2656,7 +2658,8 @@ public Operator visitShuffleSink(ShuffleSinkNode node, LocalExecutionPlanContext localInstanceId.toThrift(), node.getPlanNodeId().getId(), context.getInstanceContext()); - context.setISink(sinkHandle); + sinkHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + context.getDriverContext().setSink(sinkHandle); return new ShuffleHelperOperator(operatorContext, children, downStreamChannelIndex, sinkHandle); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java index f5a716a72b8f..47c31964a3ff 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/SharedExchangeMemoryEligibility.java @@ -20,8 +20,9 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; -import org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.ExchangeNode; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode; import java.util.ArrayDeque; import java.util.Deque; @@ -37,20 +38,16 @@ static boolean isEligible(PlanNode root) { nodes.push(root); while (!nodes.isEmpty()) { PlanNode node = nodes.pop(); - if (node instanceof ExchangeNode - || node - instanceof - org.apache.iotdb.db.queryengine.plan.relational.planner.node.ExchangeNode) { + if (node instanceof ExchangeNode) { // An upstream plan, if still attached, belongs to a different FI. continue; } + // TopK's input ordering assumptions remain unchanged, including its single-child form. + if (node instanceof TopKNode) { + return false; + } // Core relational nodes need not implement the serialization-only PlanNode.getType(). - boolean anyReady = - node instanceof CollectNode - || node instanceof UnionNode - || node - instanceof - org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.CollectNode; + boolean anyReady = node instanceof CollectNode || node instanceof UnionNode; hasCollect |= anyReady; List children = node.getChildren(); // Unknown multi-input operators are not implicitly safe. In particular, sequential sinks diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java index d7461431824e..016ee576584b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java @@ -630,27 +630,11 @@ public List visitTopK(TopKNode node, PlanContext context) { if (canTopKEliminated(node.getOrderingScheme(), node.getCount(), childrenNodes.get(0))) { return childrenNodes; } - TopKNode physical = - new TopKNode( - node.getPlanNodeId(), - childrenNodes, - node.getOrderingScheme(), - node.getCount(), - node.getOutputSymbols(), - canSortEliminated( - node.getOrderingScheme(), - nodeOrderingMap.get(childrenNodes.get(0).getPlanNodeId()))); - physical.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); - return Collections.singletonList(physical); + node.setChildren(Collections.singletonList(childrenNodes.get(0))); + return Collections.singletonList(node); } - TopKNode newTopKNode = - new TopKNode( - node.getPlanNodeId(), - node.getOrderingScheme(), - node.getCount(), - node.getOutputSymbols(), - true); + TopKNode newTopKNode = (TopKNode) node.clone(); // root TopK is not a runtime filter producer. newTopKNode.setTopKRuntimeFilterSourceId(null); for (PlanNode child : childrenNodes) { @@ -665,8 +649,7 @@ public List visitTopK(TopKNode node, PlanContext context) { node.getOrderingScheme(), node.getCount(), node.getOutputSymbols(), - canSortEliminated( - node.getOrderingScheme(), nodeOrderingMap.get(child.getPlanNodeId()))); + node.isChildrenDataInOrder()); regionTopK.setTopKRuntimeFilterSourceId(node.getTopKRuntimeFilterSourceId()); newChild = regionTopK; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index ebcc275717f6..49623ab9e45e 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -21,6 +21,7 @@ import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator; import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; import org.apache.iotdb.commons.memory.MemoryManager; @@ -28,6 +29,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.UnionNode; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; @@ -45,8 +47,10 @@ import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; +import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; import com.google.common.collect.ArrayListMultimap; +import com.google.common.util.concurrent.SettableFuture; import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.read.common.block.TsBlock; @@ -63,6 +67,8 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.same; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -99,7 +105,6 @@ private LocalExecutionPlanContext newContext() { mock(MemoryReservationManager.class)); LocalExecutionPlanContext result = new LocalExecutionPlanContext(new TypeProvider(), instance, (DataNodeQueryContext) null); - result.addPipelineDriverFactory(mock(Operator.class), result.getDriverContext(), 0); result.setExchangeSumNum(4); return result; } @@ -117,6 +122,9 @@ private static PlanNode node(PlanNodeType type, PlanNode... children) { case TABLE_COLLECT_NODE: nodeClass = CollectNode.class; break; + case TABLE_TOPK_NODE: + nodeClass = TopKNode.class; + break; case TREE_COLLECT: nodeClass = org.apache.iotdb.db.queryengine.plan.planner.plan.node.process.CollectNode.class; @@ -139,6 +147,21 @@ private static PlanNode collect() { Collections.emptyList()); } + private CollectOperator generateTableCollect(PlanNode plan) { + Operator first = mock(Operator.class); + Operator second = mock(Operator.class); + when(first.isBlocked()).thenReturn(SettableFuture.create()); + doReturn(Operator.NOT_BLOCKED).when(second).isBlocked(); + when(first.calculateMaxPeekMemoryWithCounter()).thenReturn(100L); + when(first.calculateRetainedSizeAfterCallingNext()).thenReturn(30L); + when(second.calculateMaxPeekMemoryWithCounter()).thenReturn(80L); + when(second.calculateRetainedSizeAfterCallingNext()).thenReturn(50L); + doReturn(first).when(plan.getChildren().get(0)).accept(any(), same(context)); + doReturn(second).when(plan.getChildren().get(1)).accept(any(), same(context)); + return (CollectOperator) + plan.accept(new DataNodeTableOperatorGenerator(mock(Metadata.class)), context); + } + private ISourceHandle addSource(LocalExecutionPlanContext target) { ISourceHandle handle = mock(SourceHandle.class); ExchangeOperator operator = mock(ExchangeOperator.class); @@ -296,14 +319,82 @@ public void testMappingCollectUsesSharedBudget() { } @Test - public void testTreeCollectUsesSharedBudget() { + public void testTreePlanningRetainsLegacySourceQuota() { + ISourceHandle source = addSource(context); + context.setMaxBytesOneHandleCanReserve(); + verify(source).setMaxBytesCanReserve(256); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TREE_COLLECT, + node(PlanNodeType.EXCHANGE), + node(PlanNodeType.EXCHANGE)))); + assertAdmissionIsShared(false); + } + + @Test + public void testTableCollectEnablesReadyConsumptionOnlyForEligibleFragment() { + PlanNode plan = collect(); + CollectOperator collect = generateTableCollect(plan); addSource(context); - context.setExchangeMemoryLimits( - node(PlanNodeType.TREE_COLLECT, node(PlanNodeType.EXCHANGE), node(PlanNodeType.EXCHANGE)), - pool); + context.setExchangeMemoryLimits(plan, pool); + assertTrue(collect.isBlocked().isDone()); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); assertAdmissionIsShared(true); } + @Test + public void testTableUnionEnablesReadyConsumptionOnlyForEligibleFragment() { + PlanNode union = + new UnionNode( + new PlanNodeId("union"), + Arrays.asList( + node(PlanNodeType.TABLE_EXCHANGE_NODE), node(PlanNodeType.TABLE_EXCHANGE_NODE)), + ArrayListMultimap.create(), + Collections.emptyList()); + CollectOperator collect = generateTableCollect(union); + addSource(context); + context.setExchangeMemoryLimits(union, pool); + assertTrue(collect.isBlocked().isDone()); + assertEquals(80, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(150, collect.calculateMaxPeekMemory()); + assertAdmissionIsShared(true); + } + + @Test + public void testSingleChildTopKRetainsLegacyBudgetAndConsumption() { + PlanNode child = collect(); + CollectOperator collect = generateTableCollect(child); + ISourceHandle source = addSource(context); + PlanNode topK = node(PlanNodeType.TABLE_TOPK_NODE, child); + context.setExchangeMemoryLimits(topK, pool); + verify(source).setMaxBytesCanReserve(256); + assertFalse(collect.isBlocked().isDone()); + assertAdmissionIsShared(false); + assertFalse( + SharedExchangeMemoryEligibility.isEligible( + node(PlanNodeType.TABLE_COLLECT_NODE, topK, collect()))); + assertTrue( + SharedExchangeMemoryEligibility.isEligible( + node( + PlanNodeType.TABLE_COLLECT_NODE, + node(PlanNodeType.TABLE_EXCHANGE_NODE, topK), + node(PlanNodeType.TABLE_EXCHANGE_NODE)))); + } + + @Test + public void testMultiplePipelinesRetainSequentialCollectConsumption() { + PlanNode plan = collect(); + CollectOperator collect = generateTableCollect(plan); + LocalExecutionPlanContext child = context.createSubContext(); + child.addPipelineDriverFactory(mock(Operator.class), child.getDriverContext(), 0); + addSource(context); + context.setExchangeMemoryLimits(plan, pool); + assertFalse(collect.isBlocked().isDone()); + assertAdmissionIsShared(false); + } + @Test public void testNestedDependencyDisablesSharingForEntireFragment() { for (PlanNodeType dependency : @@ -393,7 +484,7 @@ public void testInternalPipelineDisablesSharingForEntireFragment() { @Test public void testPipelineSchedulingDependencyDisablesSharing() { - context.getCurrentPipelineDriverFactory().setDependencyPipeline(0); + context.getDriverContext().setDependencyDriverIndex(0); context.setExchangeMemoryLimits(collect(), pool); assertAdmissionIsShared(false); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java index 5bae986bf476..aa20dd0005bc 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/CollectInputPropertyTest.java @@ -25,7 +25,6 @@ import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.GroupNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.SortNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.StreamSortNode; -import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.TopKNode; import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.WindowNode; import org.junit.Test; @@ -148,24 +147,6 @@ public void testUnionFinalAggregationDoesNotInheritBranchGrouping() { assertTrue(foundFinalCollect); } - @Test - public void testTopKAboveGlobalWindowDoesNotAssumeOrderedInput() { - PlanNode root = - plan( - "SELECT tag1, tag2, tag3, s1, COUNT(*) OVER () AS n FROM table1 " - + "ORDER BY s1 DESC, tag1, tag2, tag3 LIMIT 5"); - List topKs = nodes(root, TopKNode.class); - assertFalse(topKs.isEmpty()); - boolean foundUnorderedInput = false; - for (TopKNode topK : topKs) { - if (!nodes(topK.getChildren().get(0), CollectNode.class).isEmpty()) { - foundUnorderedInput = true; - assertFalse(topK.isChildrenDataInOrder()); - } - } - assertTrue(foundUnorderedInput); - } - @Test public void testOrderedScanBranchesKeepStreamingAggregation() { PlanNode root = From 5d24b4c85895944afcbd9dc29f46069a55d79ce6 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Fri, 16 Oct 2026 01:59:20 +0800 Subject: [PATCH 13/15] Address ready Collect review findings and deterministic replica checks Cancel pending reservations before FI deregistration and keep legacy accounting independent per plan node. Guard late Collect mode changes, cover planner entry points, and order pipe verification queries by their complete primary key. --- .../it/dual/tablemodel/TableModelUtils.java | 6 +- .../it/query/recent/IoTDBCollectFanInIT.java | 21 ++ .../apache/iotdb/calc/i18n/CalcMessages.java | 3 + .../apache/iotdb/calc/i18n/CalcMessages.java | 3 + .../operator/process/CollectOperator.java | 10 + .../operator/process/CollectOperatorTest.java | 114 ++++++++ .../execution/memory/MemoryPool.java | 33 ++- .../execution/memory/MemoryPoolTest.java | 156 ++++++++++- .../LocalExecutionPlannerBudgetTest.java | 245 ++++++++++++++++++ 9 files changed, 580 insertions(+), 11 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java index 2aa5b49f331f..700d552f0c9a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java @@ -366,7 +366,11 @@ public static String generateHeaderResults() { } public static String getQuerySql(final String table) { - return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + table; + // ClusterTestResultSet compares replicas row by row. All tags plus time form a unique key, + // so require the same row order even when Collect consumes inputs in readiness order. + return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + + table + + " order by s0,s1,s2,s3,time"; } public static String getQueryCountSql(final String table) { diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java index 356908b5affc..72e75a9cfb9c 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java @@ -19,10 +19,12 @@ package org.apache.iotdb.relational.it.query.recent; +import org.apache.iotdb.db.it.utils.TestUtils; import org.apache.iotdb.it.env.EnvFactory; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.TableClusterIT; import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; +import org.apache.iotdb.pipe.it.dual.tablemodel.TableModelUtils; import org.junit.AfterClass; import org.junit.BeforeClass; @@ -34,6 +36,7 @@ import java.sql.ResultSet; import java.sql.Statement; import java.util.ArrayList; +import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -55,6 +58,7 @@ public static void setUp() throws Exception { EnvFactory.getEnv() .getConfig() .getCommonConfig() + .setDegreeOfParallelism(1) .setMaxTsBlockLineNumber(2) .setSortBufferSize(128 * 1024) .setDataRegionGroupExtensionPolicy("CUSTOM") @@ -110,6 +114,23 @@ public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Excepti } } + @Test + public void testPipeVerificationQueryUsesDeterministicPrimaryKeyOrder() throws Exception { + TableModelUtils.createDataBaseAndTable(EnvFactory.getEnv(), "pipe_readings", DATABASE); + TableModelUtils.insertData(DATABASE, "pipe_readings", 0, DEVICES * ROWS, EnvFactory.getEnv()); + List expected = + new ArrayList<>(TableModelUtils.generateExpectedResults(0, DEVICES * ROWS)); + Collections.sort(expected); + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE); + try (ResultSet result = + statement.executeQuery(TableModelUtils.getQuerySql("pipe_readings"))) { + TestUtils.assertResultSetEqual(result, TableModelUtils.generateHeaderResults(), expected); + } + } + } + @Test public void testLegacyTreeAndSharedTreeToTablePaths() throws Exception { try (Connection connection = EnvFactory.getEnv().getConnection(); diff --git a/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java b/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java index 750def8c4a3b..33976ed731f2 100644 --- a/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java +++ b/iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java @@ -582,4 +582,7 @@ private CalcMessages() {} EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 = "Aggregate function [%s] requires %d arguments, but got %d"; + public static final String + EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 = + "Ready input consumption must be enabled before Collect execution starts"; } diff --git a/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java b/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java index 06ac8d62cbbf..19a56413d8b0 100644 --- a/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java +++ b/iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java @@ -558,4 +558,7 @@ private CalcMessages() {} public static final String EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 = "聚合函数 [%s] 需要 %d 个参数，但实际传入 %d 个"; + public static final String + EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 = + "必须在 Collect 开始执行前启用就绪输入消费"; } diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java index c6ee9577efe6..e94139ff0cdc 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java @@ -21,6 +21,7 @@ import org.apache.iotdb.calc.execution.operator.CommonOperatorContext; import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.i18n.CalcMessages; import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper; import org.apache.iotdb.commons.utils.TestOnly; @@ -30,6 +31,7 @@ import java.util.List; +import static com.google.common.base.Preconditions.checkState; import static com.google.common.util.concurrent.Futures.getDone; import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; @@ -41,6 +43,7 @@ public class CollectOperator implements ProcessOperator { protected final List children; private AnyChildBlocked anyChildBlocked; private boolean inited; + private boolean executionStarted; protected int currentIndex; private ListenableFuture selectionFuture; private ListenableFuture exposedFuture; @@ -58,6 +61,10 @@ public CollectOperator(CommonOperatorContext operatorContext, List chi */ public void enableReadyInputConsumption() { if (anyChildBlocked == null) { + checkState( + !executionStarted, + CalcMessages + .EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4); anyChildBlocked = new AnyChildBlocked(children); } } @@ -71,6 +78,7 @@ public boolean hasNext() { @Override public TsBlock next() throws Exception { + executionStarted = true; if (anyChildBlocked == null) { if (children.get(currentIndex).hasNextWithTimer()) { TsBlock block = children.get(currentIndex).nextWithTimer(); @@ -137,6 +145,7 @@ protected void closeCurrentChild(int index) throws Exception { @Override public ListenableFuture isBlocked() { + executionStarted = true; if (anyChildBlocked == null) { if (!inited) { inited = true; @@ -170,6 +179,7 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { + executionStarted = true; if (anyChildBlocked == null) { for (int i = currentIndex; i < children.size(); i++) { if (children.get(i) != null) { diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java index 6ab0b3a35f92..ac05610712aa 100644 --- a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -55,6 +55,7 @@ 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; @@ -425,6 +426,119 @@ public void testDefaultMappingCollectRetainsSequentialConsumption() throws Excep collect.close(); } + @Test + public void testLegacyInitializesOtherInputsOnceAndFinishesAtLastChild() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + for (int i = 0; i < 10; i++) { + assertFalse(collect.isBlocked().isDone()); + } + assertEquals(11, a.waitCalls); + assertEquals(1, b.waitCalls); + assertEquals(0, b.advances); + a.ready(); + assertSame(a.output, collect.next()); + a.hasNext = false; + assertNull(collect.next()); + assertTrue(collect.hasNext()); + assertFalse(collect.isFinished()); + for (int i = 0; i < 10; i++) { + assertTrue(collect.isBlocked().isDone()); + } + assertEquals(11, a.waitCalls); + assertEquals(11, b.waitCalls); + assertSame(b.output, collect.next()); + b.hasNext = false; + assertNull(collect.next()); + assertFalse(collect.hasNext()); + assertTrue(collect.isFinished()); + assertTrue(collect.isBlocked().isDone()); + collect.close(); + assertEquals(1, a.closes); + assertEquals(1, b.closes); + } + + @Test + public void testLegacyMemoryUsesLargestChildPeak() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.peak = 100; + a.retained = 30; + a.returnSize = 10; + b.peak = 80; + b.retained = 50; + b.returnSize = 20; + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertEquals(100, collect.calculateMaxPeekMemory()); + assertEquals(0, collect.calculateRetainedSizeAfterCallingNext()); + assertEquals(20, collect.calculateMaxReturnSize()); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyInputIsRetired() throws Exception { + Child a = new Child(block(1)); + Child b = new Child(block(2)); + a.hasNext = false; + b.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b))); + assertNull(collect.next()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + assertSame(b.output, collect.next()); + b.hasNext = false; + assertNull(collect.next()); + assertTrue(collect.isFinished()); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyReadinessObservation() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + assertFalse(collect.isBlocked().isDone()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyNextWithoutIsBlocked() throws Exception { + Child child = new Child(block(1)); + child.ready(); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + assertSame(child.output, collect.next()); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + collect.close(); + } + + @Test + public void testCannotEnableReadyConsumptionAfterLegacyClose() throws Exception { + Child child = new Child(block(1)); + CollectOperator collect = + new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child))); + collect.close(); + assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption); + assertEquals(1, child.closes); + } + + @Test + public void testEnablingReadyConsumptionIsIdempotent() throws Exception { + Child child = new Child(block(1)); + child.ready(); + CollectOperator collect = collect(child); + collect.enableReadyInputConsumption(); + ListenableFuture selection = collect.isBlocked(); + collect.enableReadyInputConsumption(); + assertSame(selection, collect.isBlocked()); + assertSame(child.output, collect.next()); + collect.close(); + } + @Test public void testMappingUsesActualChildIdentity() throws Exception { Child a = new Child(block(10, 11)); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 219372c9bd34..6a98ca7c4828 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -174,7 +174,7 @@ private static class FragmentMemoryReservation { private final AtomicLong inputUsedBytes = new AtomicLong(); private final AtomicLong outputUsedBytes = new AtomicLong(); // Enabled by local planning before drivers start; request limits are supplied by handles. - private boolean sharedBudget; + private volatile boolean sharedBudget; private AtomicLong getUsedBytes(MemorySide side) { return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes; @@ -299,6 +299,18 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( // forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed. return; } + // Pending requests must not outlive their FI registration. Cancel before removing the + // account so another query's free() cannot retry a request against a missing FI. Do not + // hold the registration lock while notifying cancellation listeners. + Iterator> pending = memoryReservationFutures.iterator(); + while (pending.hasNext()) { + MemoryReservationFuture future = pending.next(); + if (queryId.equals(future.getQueryId()) + && fragmentInstanceId.equals(future.getFragmentInstanceId())) { + tryCancel(future); + pending.remove(); + } + } synchronized (queryMemoryReservations) { queryRelatedMemory.remove(fragmentInstanceId); // Another fragment may have recreated the query entry while this cleanup waited for the @@ -451,10 +463,11 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l DataNodeQueryMessages .QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2); } - synchronized (fragment) { - PlanNodeReservation reservation = - requireReservation( - fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + Object accountingLock = fragment.sharedBudget ? fragment : reservation; + synchronized (accountingLock) { if (reservation.bytesReserved < bytes) { throw new IllegalArgumentException( DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED); @@ -540,10 +553,12 @@ public boolean tryReserve( long maxBytesCanReserve) { FragmentMemoryReservation fragment = getFragmentReservation(queryId, fragmentInstanceId, planNodeId); - synchronized (fragment) { - PlanNodeReservation reservation = - requireReservation( - fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + PlanNodeReservation reservation = + requireReservation( + fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId); + // Legacy quotas are independent per plan node. Only shared quotas require the FI lock. + Object accountingLock = fragment.sharedBudget ? fragment : reservation; + synchronized (accountingLock) { long usedBytes = fragment.sharedBudget ? fragment.getUsedBytes(reservation.side).get() diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index ed19981b400d..25957b97897d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -180,6 +180,85 @@ public void testForceDeregisterReportsMemoryLeak() { Assert.assertEquals(256L, pool.getReservedBytes()); } + @Test + public void testDeregisterCancelsPendingReservationsBeforeUnrelatedFree() { + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "b", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 512, 512)); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "b", 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + Assert.assertFalse(pending.getFuture().isDone()); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + + Assert.assertTrue(pending.getFuture().isCancelled()); + pool.free("healthy", "fi", "a", 512); + pool.free("healthy", "fi", "b", 512); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testForceDeregisterCancelsPendingReservationsEvenWhenReportingLeak() { + Assert.assertTrue(pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 256, 512)); + + Assert.assertThrows( + MemoryLeakException.class, + () -> + pool.deRegisterFragmentInstanceFromQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, true)); + + Assert.assertTrue(pending.getFuture().isCancelled()); + pool.free("healthy", "fi", "a", 256); + // Forced cleanup still reports and retains the original leaked allocation. + Assert.assertEquals(512, pool.getReservedBytes()); + } + + @Test + public void testNonForcedDeregisterKeepsPendingReservationWhileMemoryIsLive() { + Assert.assertTrue(pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512)); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + + Assert.assertFalse(pending.getFuture().isDone()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512); + Assert.assertTrue(pending.getFuture().isDone()); + Assert.assertFalse(pending.getFuture().isCancelled()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256); + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, false); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testDeregisterDoesNotCancelOtherFragmentPendingReservation() { + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "a", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("healthy", "fi", "b", MemorySide.INPUT); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "a", 512, 512)); + Assert.assertTrue(pool.tryReserve("healthy", "fi", "b", 512, 512)); + pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, "other", PLAN_NODE_ID, MemorySide.INPUT); + MemoryReservationResult pending = + pool.reserveWithPriority(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 256, 512, false); + MemoryReservationResult other = + pool.reserveWithPriority(QUERY_ID, "other", PLAN_NODE_ID, 256, 512, false); + + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + + Assert.assertTrue(pending.getFuture().isCancelled()); + Assert.assertFalse(other.getFuture().isDone()); + pool.free("healthy", "fi", "a", 512); + Assert.assertTrue(other.getFuture().isDone()); + Assert.assertFalse(other.getFuture().isCancelled()); + pool.free(QUERY_ID, "other", PLAN_NODE_ID, 256); + pool.free("healthy", "fi", "b", 512); + Assert.assertEquals(0, pool.getReservedBytes()); + } + @Test public void testTryReserve() { @@ -891,7 +970,19 @@ public void testSharedAndLegacyFragmentsDoNotChangeEachOthersPolicy() { @Test public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { - pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(true); + } + + @Test + public void testLegacyConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Exception { + assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(false); + } + + private void assertConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits(boolean shared) + throws Exception { + if (shared) { + pool.registerFragmentInstanceSharedBudget(QUERY_ID, FRAGMENT_INSTANCE_ID); + } for (int i = 0; i < 3; i++) { String fi = "other" + i; pool.registerPlanNodeIdToQueryMemoryMap(QUERY_ID, fi, PLAN_NODE_ID, MemorySide.OUTPUT); @@ -938,4 +1029,67 @@ public void testConcurrentAdmissionAlwaysGrantsOneRequestWhenOneFits() throws Ex for (int i = 0; i < 3; i++) pool.free(QUERY_ID, "other" + i, PLAN_NODE_ID, 256); Assert.assertEquals(0, pool.getReservedBytes()); } + + @Test + public void testConcurrentLegacyReservationsKeepIndependentQuotasAndAccurateSideTotals() + throws Exception { + for (boolean samePlanNode : new boolean[] {true, false}) { + MemoryPool legacyPool = new MemoryPool("legacy", new MemoryManager(1024), 512); + ExecutorService executor = Executors.newFixedThreadPool(16); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(16); + CountDownLatch release = new CountDownLatch(1); + List> tasks = new ArrayList<>(); + try { + for (int i = 0; i < 16; i++) { + String id = samePlanNode ? "same" : "pipeline-" + i; + legacyPool.registerPlanNodeIdToQueryMemoryMap( + QUERY_ID, FRAGMENT_INSTANCE_ID, id, MemorySide.INPUT); + tasks.add( + executor.submit( + () -> { + start.await(); + boolean reserved = + legacyPool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32, 128); + attempted.countDown(); + release.await(); + if (reserved) { + legacyPool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, id, 32); + } + return null; + })); + } + start.countDown(); + Assert.assertTrue(attempted.await(10, TimeUnit.SECONDS)); + long expected = samePlanNode ? 128 : 512; + Assert.assertEquals(expected, legacyPool.getReservedBytes()); + Assert.assertEquals( + expected, + legacyPool.getFragmentMemoryReservedBytes( + QUERY_ID, FRAGMENT_INSTANCE_ID, MemorySide.INPUT)); + if (samePlanNode) { + Assert.assertEquals( + 128, + legacyPool.getPlanNodeMemoryReservedBytes(QUERY_ID, FRAGMENT_INSTANCE_ID, "same")); + } else { + for (int i = 0; i < 16; i++) { + Assert.assertEquals( + 32, + legacyPool.getPlanNodeMemoryReservedBytes( + QUERY_ID, FRAGMENT_INSTANCE_ID, "pipeline-" + i)); + } + } + release.countDown(); + for (Future task : tasks) { + task.get(10, TimeUnit.SECONDS); + } + Assert.assertEquals(0, legacyPool.getReservedBytes()); + Assert.assertEquals(0, legacyPool.getQueryMemoryReservedBytes(QUERY_ID)); + } finally { + start.countDown(); + release.countDown(); + executor.shutdownNow(); + } + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java new file mode 100644 index 000000000000..ac32c831c43e --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java @@ -0,0 +1,245 @@ +/* + * 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.planner; + +import org.apache.iotdb.calc.execution.operator.Operator; +import org.apache.iotdb.calc.execution.operator.process.CollectOperator; +import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.node.CollectNode; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; +import org.apache.iotdb.db.queryengine.execution.exchange.source.SourceHandle; +import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; +import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; +import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; +import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; +import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; + +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class LocalExecutionPlannerBudgetTest { + private long previousBudget; + private int previousBlockSize; + private boolean previousEstimation; + private int previousDataNodeId; + + @Before + public void setUp() { + previousDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0); + previousBudget = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + previousBlockSize = TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + previousEstimation = + IoTDBDescriptor.getInstance().getMemoryConfig().isEnableQueryMemoryEstimation(); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(1024); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(64); + IoTDBDescriptor.getInstance().getMemoryConfig().setEnableQueryMemoryEstimation(true); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId); + IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(previousBudget); + TSFileDescriptor.getInstance().getConfig().setMaxTsBlockSizeInBytes(previousBlockSize); + IoTDBDescriptor.getInstance() + .getMemoryConfig() + .setEnableQueryMemoryEstimation(previousEstimation); + } + + @Test + public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() throws Exception { + LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); + MemoryPool pool = + MPPDataExchangeService.getInstance() + .getMPPDataExchangeManager() + .getLocalMemoryManager() + .getQueryPool(); + FragmentInstanceContext instance = instance("table_planner_budget", SqlDialect.TABLE); + PlanNode first = mock(PlanNode.class); + PlanNode second = mock(PlanNode.class); + SourceHandle firstHandle = mock(SourceHandle.class); + SourceHandle secondHandle = mock(SourceHandle.class); + addInput(first, firstHandle, false, 100, 30); + addInput(second, secondHandle, true, 80, 50); + CollectNode plan = + spy( + new CollectNode( + new PlanNodeId("collect"), Arrays.asList(first, second), Collections.emptyList())); + doReturn(PlanNodeType.TABLE_COLLECT_NODE).when(plan).getType(); + when(first.getChildren()).thenReturn(Collections.emptyList()); + when(second.getChildren()).thenReturn(Collections.emptyList()); + long freeBefore = planner.getFreeMemoryForOperators(); + try { + List factories = + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + assertEquals(1, factories.size()); + CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); + // Check the actual reservation, not only the formula after planning has completed. + assertEquals(150 + collect.ramBytesUsed(), freeBefore - planner.getFreeMemoryForOperators()); + assertEquals(1024, instance.getInputMaxBytes()); + assertEquals(0, instance.getOutputMaxBytes()); + assertTrue(collect.isBlocked().isDone()); + verify(firstHandle).setMaxBytesCanReserve(1024); + verify(secondHandle).setMaxBytesCanReserve(1024); + + String query = instance.getId().getQueryId().getId(); + String fi = instance.getId().getFullId(); + pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "first", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "second", MemorySide.INPUT); + assertTrue( + pool.reserveWithPriority(query, fi, "first", 1024, 1024, false).isReserveSuccess()); + MemoryPool.MemoryReservationResult pending = + pool.reserveWithPriority(query, fi, "second", 64, 1024, false); + assertFalse(pending.getFuture().isDone()); + pool.free(query, fi, "first", 1024); + assertTrue(pending.getFuture().isDone()); + pool.free(query, fi, "second", 64); + collect.close(); + } finally { + pool.deRegisterFragmentInstanceFromQueryMemoryMap( + instance.getId().getQueryId().getId(), instance.getId().getFullId(), true); + long reserved = freeBefore - planner.getFreeMemoryForOperators(); + if (reserved > 0) { + planner.releaseToFreeMemoryForOperators(reserved); + } + } + } + + @Test + public void testNullSessionRetainsTreeGeneratorAndIndependentInputQuotas() throws Exception { + FragmentInstanceContext instance = instance("null_session_planner_budget", null); + PlanNode plan = mock(PlanNode.class); + when(plan.getType()).thenReturn(PlanNodeType.TREE_COLLECT); + when(plan.getPlanNodeId()).thenReturn(new PlanNodeId("collect")); + SourceHandle firstHandle = mock(SourceHandle.class); + SourceHandle secondHandle = mock(SourceHandle.class); + Operator first = input(false, 100, 30); + Operator second = input(true, 80, 50); + doAnswer( + invocation -> { + assertTrue(invocation.getArgument(0) instanceof OperatorTreeGenerator); + LocalExecutionPlanContext context = invocation.getArgument(1); + registerInput(context, firstHandle); + registerInput(context, secondHandle); + CollectOperator collect = + new CollectOperator( + context + .getDriverContext() + .addOperatorContext(0, new PlanNodeId("collect"), "CollectOperator"), + new java.util.ArrayList<>(Arrays.asList(first, second))); + return collect; + }) + .when(plan) + .accept(any(), any()); + LocalExecutionPlanner planner = LocalExecutionPlanner.getInstance(); + long freeBefore = planner.getFreeMemoryForOperators(); + try { + List factories = + planner.plan(plan, new TypeProvider(), instance, new DataNodeQueryContext(1)); + assertEquals(1, factories.size()); + CollectOperator collect = (CollectOperator) factories.get(0).getOperation(); + assertEquals(100 + collect.ramBytesUsed(), freeBefore - planner.getFreeMemoryForOperators()); + assertFalse(collect.isBlocked().isDone()); + assertEquals(0, instance.getInputMaxBytes()); + verify(firstHandle).setMaxBytesCanReserve(512); + verify(secondHandle).setMaxBytesCanReserve(512); + collect.close(); + } finally { + long reserved = freeBefore - planner.getFreeMemoryForOperators(); + if (reserved > 0) { + planner.releaseToFreeMemoryForOperators(reserved); + } + } + } + + private static FragmentInstanceContext instance(String query, SqlDialect dialect) { + FragmentInstanceContext instance = + spy( + FragmentInstanceContext.createFragmentInstanceContext( + new FragmentInstanceId(new PlanFragmentId(new QueryId(query), 0), "0"), + mock(FragmentInstanceStateMachine.class), + mock(MemoryReservationManager.class))); + SessionInfo session = null; + if (dialect != null) { + session = mock(SessionInfo.class); + when(session.getSqlDialect()).thenReturn(dialect); + } + doReturn(session).when(instance).getSessionInfo(); + return instance; + } + + private static void addInput( + PlanNode node, SourceHandle handle, boolean ready, long peak, long retained) { + Operator operator = input(ready, peak, retained); + doAnswer( + invocation -> { + LocalExecutionPlanContext context = invocation.getArgument(1); + registerInput(context, handle); + return operator; + }) + .when(node) + .accept(any(), any()); + } + + private static void registerInput(LocalExecutionPlanContext context, SourceHandle handle) { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(handle); + context.addExchangeOperator(exchange); + context.addExchangeSumNum(1); + } + + private static Operator input(boolean ready, long peak, long retained) { + Operator operator = mock(Operator.class); + doReturn(ready ? Operator.NOT_BLOCKED : SettableFuture.create()).when(operator).isBlocked(); + when(operator.calculateMaxPeekMemoryWithCounter()).thenReturn(peak); + when(operator.calculateRetainedSizeAfterCallingNext()).thenReturn(retained); + return operator; + } +} From 54b0c06e89bde2a97173468866c6d99b3676a805 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Fri, 16 Oct 2026 09:56:33 +0800 Subject: [PATCH 14/15] Check receiver snapshots independently on each node --- .../it/triple/IoTDBShowReceiversTripleIT.java | 103 ++++++++++++------ 1 file changed, 67 insertions(+), 36 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java index 0712387775b9..5d5313d9f891 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java @@ -54,6 +54,8 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.apache.iotdb.db.it.utils.CollectQueryAssertions.getUnorderedResultSets; + @RunWith(IoTDBTestRunner.class) @Category({MultiClusterIT3.class}) public class IoTDBShowReceiversTripleIT { @@ -214,37 +216,66 @@ private boolean hasReceiversOnDifferentDataNodes( final String secondPipeName, final int secondReceiverDataNodeId) throws SQLException { - boolean hasFirstPipe = false; - boolean hasSecondPipe = false; - final Set receiverDataNodeIds = new HashSet<>(); try (final Connection connection = receiverEnv.getConnection( SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); final Statement statement = connection.createStatement(); final ResultSet resultSet = statement.executeQuery(sql)) { - while (resultSet.next()) { - final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); - if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) - || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) - || pipeIds == null) { - continue; - } - final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); - if (receiverDataNodeId == firstReceiverDataNodeId - && pipeIds.contains(firstPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasFirstPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); - } - if (receiverDataNodeId == secondReceiverDataNodeId - && pipeIds.contains(secondPipeName + "@") - && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 - && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { - hasSecondPipe = true; - receiverDataNodeIds.add(receiverDataNodeId); + // Each node must see both senders, but unordered table queries may return different row + // orders. + for (final ResultSet nodeResult : + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) + ? getUnorderedResultSets(sql, resultSet) + : Collections.singletonList(resultSet)) { + if (!hasExpectedReceivers( + nodeResult, + sqlDialect, + firstPipeName, + firstReceiverDataNodeId, + secondPipeName, + secondReceiverDataNodeId)) { + return false; } } + return true; + } + } + + private static boolean hasExpectedReceivers( + final ResultSet resultSet, + final String sqlDialect, + final String firstPipeName, + final int firstReceiverDataNodeId, + final String secondPipeName, + final int secondReceiverDataNodeId) + throws SQLException { + boolean hasFirstPipe = false; + boolean hasSecondPipe = false; + final Set receiverDataNodeIds = new HashSet<>(); + while (resultSet.next()) { + final String pipeIds = getString(resultSet, sqlDialect, "PipeIDs", "pipe_ids"); + if (!"DataNode" + .equals(getString(resultSet, sqlDialect, "ReceiverNodeType", "receiver_node_type")) + || !"thrift".equals(getString(resultSet, sqlDialect, "Protocol", "protocol")) + || pipeIds == null) { + continue; + } + final int receiverDataNodeId = + getInt(resultSet, sqlDialect, "ReceiverNodeId", "receiver_node_id"); + if (receiverDataNodeId == firstReceiverDataNodeId + && pipeIds.contains(firstPipeName + "@") + && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) { + hasFirstPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } + if (receiverDataNodeId == secondReceiverDataNodeId + && pipeIds.contains(secondPipeName + "@") + && getInt(resultSet, sqlDialect, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, sqlDialect, "PipeCount", "pipe_count") >= 1) { + hasSecondPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } } return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2; } @@ -266,22 +297,22 @@ private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception } private static String getString( - final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + final ResultSet resultSet, + final String sqlDialect, + final String treeColumnName, + final String tableColumnName) throws SQLException { - try { - return resultSet.getString(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getString(tableColumnName); - } + return resultSet.getString( + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName); } private static int getInt( - final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + final ResultSet resultSet, + final String sqlDialect, + final String treeColumnName, + final String tableColumnName) throws SQLException { - try { - return resultSet.getInt(treeColumnName); - } catch (final SQLException ignored) { - return resultSet.getInt(tableColumnName); - } + return resultSet.getInt( + BaseEnv.TABLE_SQL_DIALECT.equals(sqlDialect) ? tableColumnName : treeColumnName); } } From 788a995c75ab59703f730996897de51b31b09fbe Mon Sep 17 00:00:00 2001 From: ColinLee Date: Fri, 16 Oct 2026 11:42:06 +0800 Subject: [PATCH 15/15] Fix shared FI account keys and indexed exchange waiter cleanup --- .../execution/memory/MemoryPool.java | 58 +++-- .../planner/LocalExecutionPlanContext.java | 3 +- .../execution/memory/MemoryPoolTest.java | 205 ++++++++++++++++++ .../LocalExecutionPlanContextBudgetTest.java | 62 +++++- .../LocalExecutionPlannerBudgetTest.java | 4 +- 5 files changed, 312 insertions(+), 20 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java index 6a98ca7c4828..3664a845d4f8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java @@ -37,9 +37,8 @@ import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.Queue; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; @@ -56,6 +55,8 @@ public static class MemoryReservationFuture extends AbstractFuture { private final String fragmentInstanceId; private final String planNodeId; private final long bytesToReserve; + // Assigned once by the pool before publication. Zero denotes an untracked future. + private long reservationId; /** * MemoryReservationFuture is created when an exchange handle tries to reserve memory from the @@ -173,6 +174,8 @@ private static class FragmentMemoryReservation { private final Map planNodes = new ConcurrentHashMap<>(); private final AtomicLong inputUsedBytes = new AtomicLong(); private final AtomicLong outputUsedBytes = new AtomicLong(); + private final Map> pendingReservations = + new ConcurrentHashMap<>(); // Enabled by local planning before drivers start; request limits are supplied by handles. private volatile boolean sharedBudget; @@ -185,8 +188,11 @@ private AtomicLong getUsedBytes(MemorySide side) { private final Map> queryMemoryReservations = new ConcurrentHashMap<>(); - private final Queue> memoryReservationFutures = - new ConcurrentLinkedQueue<>(); + // Sequence keys retain FIFO wakeup order and allow removal without scanning unrelated requests. + // Only blocked reservations use these indexes; successful admission keeps its existing fast path. + private final Map> memoryReservationFutures = + new ConcurrentSkipListMap<>(); + private final AtomicLong nextReservationId = new AtomicLong(); public MemoryPool(String id, MemoryManager memoryManager, long maxBytesPerFragmentInstance) { this.id = Validate.notNull(id, DataNodeQueryMessages.EXCEPTION_ID_CAN_NOT_BE_NULL_DOT_BDD2AD7D); @@ -302,13 +308,10 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap( // Pending requests must not outlive their FI registration. Cancel before removing the // account so another query's free() cannot retry a request against a missing FI. Do not // hold the registration lock while notifying cancellation listeners. - Iterator> pending = memoryReservationFutures.iterator(); - while (pending.hasNext()) { - MemoryReservationFuture future = pending.next(); - if (queryId.equals(future.getQueryId()) - && fragmentInstanceId.equals(future.getFragmentInstanceId())) { + if (fragmentRelatedMemory != null) { + for (MemoryReservationFuture future : + fragmentRelatedMemory.pendingReservations.values()) { tryCancel(future); - pending.remove(); } } synchronized (queryMemoryReservations) { @@ -392,11 +395,21 @@ public MemoryReservationResult reserveWithPriority( MemoryReservationFuture result = MemoryReservationFuture.create( queryId, fragmentInstanceId, planNodeId, bytesToReserve, maxBytesCanReserve); - memoryReservationFutures.add(result); + FragmentMemoryReservation fragment = + getFragmentReservation(queryId, fragmentInstanceId, planNodeId); + result.reservationId = nextReservationId.incrementAndGet(); + fragment.pendingReservations.put(result.reservationId, result); + memoryReservationFutures.put(result.reservationId, result); // A concurrent free may have missed this request before it was published. Retry after // publication, under the same future lock used by free and tryCancel to avoid charging twice. - if (tryCompleteReservation(result)) { - memoryReservationFutures.remove(result); + try { + if (tryCompleteReservation(result)) { + removePendingReservation(result); + } + } catch (RuntimeException | Error e) { + // A concurrently removed registration must not leave a published request behind. + tryCancel(result); + throw e; } return new MemoryReservationResult(result, false, bytesToReserve); } @@ -440,6 +453,9 @@ public synchronized long tryCancel(ListenableFuture future) { synchronized (future) { // If the future is not a MemoryReservationFuture, it must have been completed. if (future.isDone()) { + if (future instanceof MemoryReservationFuture) { + removePendingReservation((MemoryReservationFuture) future); + } return 0L; } Validate.isTrue( @@ -447,6 +463,7 @@ public synchronized long tryCancel(ListenableFuture future) { DataNodeQueryMessages.EXCEPTION_INVALID_FUTURE_TYPE_14507EF5 + future.getClass().getSimpleName()); future.cancel(true); + removePendingReservation((MemoryReservationFuture) future); } return ((MemoryReservationFuture) future).getBytesToReserve(); } @@ -481,15 +498,26 @@ public void free(String queryId, String fragmentInstanceId, String planNodeId, l if (memoryReservationFutures.isEmpty()) { return; } - Iterator> iterator = memoryReservationFutures.iterator(); + Iterator> iterator = memoryReservationFutures.values().iterator(); while (iterator.hasNext()) { MemoryReservationFuture future = iterator.next(); if (tryCompleteReservation(future)) { - iterator.remove(); + removePendingReservation(future); } } } + private void removePendingReservation(MemoryReservationFuture future) { + memoryReservationFutures.remove(future.reservationId, future); + Map fragments = + queryMemoryReservations.get(future.getQueryId()); + FragmentMemoryReservation fragment = + fragments == null ? null : fragments.get(future.getFragmentInstanceId()); + if (fragment != null) { + fragment.pendingReservations.remove(future.reservationId, future); + } + } + private boolean tryCompleteReservation(MemoryReservationFuture future) { synchronized (future) { if (future.isDone()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java index 9e5cd2e0ea00..055bc4b19f98 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContext.java @@ -324,7 +324,8 @@ public void setExchangeMemoryLimits(PlanNode plan, MemoryPool memoryPool) { sharedBudget ? inputMaxBytes : 0, sharedBudget ? outputMaxBytes : 0); if (sharedBudget) { memoryPool.registerFragmentInstanceSharedBudget( - instanceContext.getId().getQueryId().getId(), instanceContext.getId().getFullId()); + instanceContext.getId().getQueryId().getId(), + instanceContext.getId().getFragmentInstanceId()); collectOperators.forEach(CollectOperator::enableReadyInputConsumption); LOGGER.debug( DataNodeQueryMessages.LOG_SHARED_EXCHANGE_BUDGET_FOR_FI_ARG_INPUT_ARG_OUTPUT_ARG_98AE1109, diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java index 25957b97897d..d2b901d7f235 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPoolTest.java @@ -32,8 +32,15 @@ import java.lang.management.ManagementFactory; import java.lang.management.ThreadInfo; import java.lang.reflect.Field; +import java.util.AbstractCollection; import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; import java.util.List; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -41,6 +48,7 @@ import java.util.concurrent.FutureTask; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.awaitility.Awaitility.await; @@ -259,6 +267,203 @@ public void testDeregisterDoesNotCancelOtherFragmentPendingReservation() { Assert.assertEquals(0, pool.getReservedBytes()); } + @Test + public void testOwnerCancellationAndDeregistrationDoNotScanUnrelatedWaiters() throws Exception { + assertDeregistrationDoesNotScanUnrelatedWaiters(true); + } + + @Test + public void testDeregistrationCancelsOnlyIndexedWaiters() throws Exception { + assertDeregistrationDoesNotScanUnrelatedWaiters(false); + } + + private void assertDeregistrationDoesNotScanUnrelatedWaiters(boolean ownerCancels) + throws Exception { + AtomicInteger visits = countGlobalPendingVisits(); + pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512); + pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT); + pool.tryReserve("holder", "fi", "p", 512, 512); + int count = 1024; + List victims = new ArrayList<>(); + List unrelated = new ArrayList<>(); + for (int i = 0; i < count; i++) { + String fi = Integer.toString(i); + pool.registerPlanNodeIdToQueryMemoryMap("victim", fi, "p", MemorySide.INPUT); + pool.registerPlanNodeIdToQueryMemoryMap("unrelated", fi, "p", MemorySide.INPUT); + victims.add(pool.reserveWithPriority("victim", fi, "p", 8, 512, false)); + unrelated.add(pool.reserveWithPriority("unrelated", fi, "p", 8, 512, false)); + } + for (int i = 0; i < count; i++) { + if (ownerCancels) { + Assert.assertEquals(8, pool.tryCancel(victims.get(i).getFuture())); + } + pool.deRegisterFragmentInstanceFromQueryMemoryMap("victim", Integer.toString(i), true); + Assert.assertTrue(victims.get(i).getFuture().isCancelled()); + Assert.assertFalse(unrelated.get(i).getFuture().isDone()); + } + // An operation-count assertion avoids timing noise and detects a global scan for every FI. + Assert.assertEquals(0, visits.get()); + Assert.assertEquals(count, pool.getMemoryReservationSize()); + for (int i = 0; i < count; i++) { + Assert.assertEquals(8, pool.tryCancel(unrelated.get(i).getFuture())); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("unrelated", Integer.toString(i), true); + } + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512); + pool.free("holder", "fi", "p", 512); + pool.deRegisterFragmentInstanceFromQueryMemoryMap(QUERY_ID, FRAGMENT_INSTANCE_ID, true); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true); + Assert.assertEquals(0, pool.getMemoryReservationSize()); + Assert.assertEquals(0, pool.getQueryMemoryReservationSize()); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @SuppressWarnings("unchecked") + private AtomicInteger countGlobalPendingVisits() throws Exception { + AtomicInteger visits = new AtomicInteger(); + Field field = MemoryPool.class.getDeclaredField("memoryReservationFutures"); + field.setAccessible(true); + Object pending = field.get(pool); + if (pending instanceof Queue) { + // Keep this branch so the regression also runs against the former queue implementation. + Queue> counted = + new ConcurrentLinkedQueue>() { + @Override + public Iterator> iterator() { + return countVisits(super.iterator(), visits); + } + }; + counted.addAll((Queue>) pending); + field.set(pool, counted); + } else { + Map> counted = + new ConcurrentSkipListMap>() { + @Override + public Collection> values() { + Collection> values = super.values(); + return new AbstractCollection>() { + @Override + public Iterator> iterator() { + return countVisits(values.iterator(), visits); + } + + @Override + public int size() { + return values.size(); + } + }; + } + }; + counted.putAll((Map>) pending); + field.set(pool, counted); + } + return visits; + } + + private static Iterator countVisits(Iterator delegate, AtomicInteger visits) { + return new Iterator() { + @Override + public boolean hasNext() { + return delegate.hasNext(); + } + + @Override + public T next() { + visits.incrementAndGet(); + return delegate.next(); + } + + @Override + public void remove() { + delegate.remove(); + } + }; + } + + @Test + public void testCancellationKeepsFifoWakeupOrder() { + pool.tryReserve(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 512, 512); + pool.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "p", MemorySide.INPUT); + pool.tryReserve("holder", "fi", "p", 512, 512); + List requests = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + String fi = Integer.toString(i); + pool.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT); + requests.add(pool.reserveWithPriority("waiting", fi, "p", 64, 512, false)); + } + Assert.assertEquals(64, pool.tryCancel(requests.get(1).getFuture())); + Assert.assertEquals(0, pool.tryCancel(requests.get(1).getFuture())); + Assert.assertEquals(2, pool.getMemoryReservationSize()); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 64); + Assert.assertTrue(requests.get(0).getFuture().isDone()); + Assert.assertFalse(requests.get(2).getFuture().isDone()); + pool.free("waiting", "0", "p", 64); + Assert.assertTrue(requests.get(2).getFuture().isDone()); + pool.free("waiting", "2", "p", 64); + pool.free(QUERY_ID, FRAGMENT_INSTANCE_ID, PLAN_NODE_ID, 448); + pool.free("holder", "fi", "p", 512); + for (int i = 0; i < 3; i++) { + pool.deRegisterFragmentInstanceFromQueryMemoryMap("waiting", Integer.toString(i), true); + } + Assert.assertEquals(0, pool.getMemoryReservationSize()); + Assert.assertEquals(0, pool.getReservedBytes()); + } + + @Test + public void testConcurrentCancellationAndGrantRemoveAllPendingIndexes() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int round = 0; round < 32; round++) { + MemoryPool concurrent = new MemoryPool("race", new MemoryManager(1024), 512); + concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "a", MemorySide.INPUT); + concurrent.registerPlanNodeIdToQueryMemoryMap("holder", "fi", "b", MemorySide.INPUT); + concurrent.tryReserve("holder", "fi", "a", 512, 512); + concurrent.tryReserve("holder", "fi", "b", 512, 512); + List requests = new ArrayList<>(); + for (int i = 0; i < 16; i++) { + String fi = Integer.toString(i); + concurrent.registerPlanNodeIdToQueryMemoryMap("waiting", fi, "p", MemorySide.INPUT); + requests.add(concurrent.reserveWithPriority("waiting", fi, "p", 16, 512, false)); + } + CountDownLatch start = new CountDownLatch(1); + Future grant = + executor.submit( + () -> { + start.await(); + concurrent.free("holder", "fi", "a", 512); + return null; + }); + Future cancel = + executor.submit( + () -> { + start.await(); + for (MemoryReservationResult request : requests) { + concurrent.tryCancel(request.getFuture()); + } + return null; + }); + start.countDown(); + grant.get(10, TimeUnit.SECONDS); + cancel.get(10, TimeUnit.SECONDS); + for (int i = 0; i < requests.size(); i++) { + Assert.assertTrue(requests.get(i).getFuture().isDone()); + if (!requests.get(i).getFuture().isCancelled()) { + concurrent.free("waiting", Integer.toString(i), "p", 16); + } + concurrent.deRegisterFragmentInstanceFromQueryMemoryMap( + "waiting", Integer.toString(i), true); + } + concurrent.free("holder", "fi", "b", 512); + concurrent.deRegisterFragmentInstanceFromQueryMemoryMap("holder", "fi", true); + Assert.assertEquals(0, concurrent.getMemoryReservationSize()); + Assert.assertEquals(0, concurrent.getQueryMemoryReservationSize()); + Assert.assertEquals(0, concurrent.getReservedBytes()); + } + } finally { + executor.shutdownNow(); + Assert.assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } + } + @Test public void testTryReserve() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java index 49623ab9e45e..2b3c847e89e4 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanContextBudgetTest.java @@ -24,6 +24,8 @@ import org.apache.iotdb.calc.execution.operator.process.CollectOperator; import org.apache.iotdb.calc.execution.operator.process.TableMergeSortOperator; import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.memory.MemoryManager; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; @@ -35,6 +37,7 @@ import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.SourceHandleListener; import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISinkHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; import org.apache.iotdb.db.queryengine.execution.exchange.source.LocalSourceHandle; @@ -43,11 +46,13 @@ import org.apache.iotdb.db.queryengine.execution.fragment.DataNodeQueryContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext; import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; +import org.apache.iotdb.db.queryengine.execution.memory.LocalMemoryManager; import org.apache.iotdb.db.queryengine.execution.memory.MemoryPool; import org.apache.iotdb.db.queryengine.execution.memory.MemorySide; import org.apache.iotdb.db.queryengine.execution.operator.source.ExchangeOperator; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata; +import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; import com.google.common.collect.ArrayListMultimap; import com.google.common.util.concurrent.SettableFuture; @@ -55,6 +60,7 @@ import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -63,6 +69,8 @@ import java.util.Arrays; import java.util.Collections; import java.util.Comparator; +import java.util.List; +import java.util.concurrent.ExecutorService; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -188,7 +196,7 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit assertEquals(0, instance.getInputMaxBytes()); assertEquals(0, instance.getOutputMaxBytes()); } - String fi = context.getFragmentInstanceId().getFullId(); + String fi = context.getFragmentInstanceId().getFragmentInstanceId(); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a" + side, side); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b" + side, side); assertTrue( @@ -206,6 +214,56 @@ private void assertAdmissionIsShared(boolean shared, MemorySide side, long limit pool.free("budget", fi, "b" + side, limit); } + @Test + public void testRealSourceHandlesShareCanonicalFragmentAccountAndCleanup() throws Exception { + FragmentInstanceId id = context.getFragmentInstanceId(); + String fi = FragmentInstanceId.createFragmentInstanceIdFromTFragmentInstanceId(id.toThrift()); + LocalMemoryManager manager = mock(LocalMemoryManager.class); + when(manager.getQueryPool()).thenReturn(pool); + SourceHandle first = realSourceHandle(id, "first", manager); + SourceHandle second = realSourceHandle(id, "second", manager); + for (SourceHandle source : Arrays.asList(first, second)) { + ExchangeOperator exchange = mock(ExchangeOperator.class); + when(exchange.getSourceHandle()).thenReturn(source); + context.addExchangeOperator(exchange); + } + try { + context.setExchangeMemoryLimits(collect(), pool); + assertEquals(1024, context.getInstanceContext().getInputMaxBytes()); + first.isBlocked(); + second.isBlocked(); + first.updatePendingDataBlockInfo(0, List.of(1024L)); + second.updatePendingDataBlockInfo(0, List.of(1024L)); + assertEquals(1024, pool.getFragmentMemoryReservedBytes("budget", fi, MemorySide.INPUT)); + assertEquals(0, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second")); + assertEquals(1, pool.getMemoryReservationSize()); + first.abort(); + assertEquals(1024, pool.getPlanNodeMemoryReservedBytes("budget", fi, "second")); + } finally { + first.abort(); + second.abort(); + pool.deRegisterFragmentInstanceFromQueryMemoryMap("budget", fi, true); + } + assertEquals(0, pool.getMemoryReservationSize()); + assertEquals(0, pool.getReservedBytes()); + assertEquals(0, pool.getQueryMemoryReservationSize()); + } + + private static SourceHandle realSourceHandle( + FragmentInstanceId id, String planNode, LocalMemoryManager manager) { + return new SourceHandle( + new TEndPoint("remote", 12345), + new TFragmentInstanceId(id.getQueryId().getId(), 1, "0"), + id.toThrift(), + planNode, + 0, + manager, + mock(ExecutorService.class), + new TsBlockSerde(), + mock(SourceHandleListener.class), + mock(IClientManager.class)); + } + @Test public void testPureCollectUsesSharedSideBudgets() { ISourceHandle source = addSource(context); @@ -515,7 +573,7 @@ public void testMergeSortGetsBothHeadsDespitePrefetchOnOneInput() throws Excepti context.setExchangeMemoryLimits( node(PlanNodeType.TABLE_MERGESORT_NODE, collect(), node(PlanNodeType.TABLE_EXCHANGE_NODE)), pool); - String fi = context.getFragmentInstanceId().getFullId(); + String fi = context.getFragmentInstanceId().getFragmentInstanceId(); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "a", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap("budget", fi, "b", MemorySide.INPUT); for (int i = 0; i < 4; i++) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java index ac32c831c43e..31c562199fd6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerBudgetTest.java @@ -130,7 +130,7 @@ public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() th verify(secondHandle).setMaxBytesCanReserve(1024); String query = instance.getId().getQueryId().getId(); - String fi = instance.getId().getFullId(); + String fi = instance.getId().getFragmentInstanceId(); pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "first", MemorySide.INPUT); pool.registerPlanNodeIdToQueryMemoryMap(query, fi, "second", MemorySide.INPUT); assertTrue( @@ -144,7 +144,7 @@ public void testTablePlanEnablesSharingBeforeMemoryEstimationAndRootFactory() th collect.close(); } finally { pool.deRegisterFragmentInstanceFromQueryMemoryMap( - instance.getId().getQueryId().getId(), instance.getId().getFullId(), true); + instance.getId().getQueryId().getId(), instance.getId().getFragmentInstanceId(), true); long reserved = freeBefore - planner.getFreeMemoryForOperators(); if (reserved > 0) { planner.releaseToFreeMemoryForOperators(reserved);