Skip to content

Commit ab11600

Browse files
ColinLeeoColinLeeo
committed
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.
1 parent ba1d7db commit ab11600

9 files changed

Lines changed: 580 additions & 11 deletions

File tree

‎integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/TableModelUtils.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -366,7 +366,11 @@ public static String generateHeaderResults() {
366366
}
367367

368368
public static String getQuerySql(final String table) {
369-
return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from " + table;
369+
// ClusterTestResultSet compares replicas row by row. All tags plus time form a unique key,
370+
// so require the same row order even when Collect consumes inputs in readiness order.
371+
return "select s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,time from "
372+
+ table
373+
+ " order by s0,s1,s2,s3,time";
370374
}
371375

372376
public static String getQueryCountSql(final String table) {

‎integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBCollectFanInIT.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@
1919

2020
package org.apache.iotdb.relational.it.query.recent;
2121

22+
import org.apache.iotdb.db.it.utils.TestUtils;
2223
import org.apache.iotdb.it.env.EnvFactory;
2324
import org.apache.iotdb.it.framework.IoTDBTestRunner;
2425
import org.apache.iotdb.itbase.category.TableClusterIT;
2526
import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
27+
import org.apache.iotdb.pipe.it.dual.tablemodel.TableModelUtils;
2628

2729
import org.junit.AfterClass;
2830
import org.junit.BeforeClass;
@@ -34,6 +36,7 @@
3436
import java.sql.ResultSet;
3537
import java.sql.Statement;
3638
import java.util.ArrayList;
39+
import java.util.Collections;
3740
import java.util.HashSet;
3841
import java.util.List;
3942
import java.util.Set;
@@ -55,6 +58,7 @@ public static void setUp() throws Exception {
5558
EnvFactory.getEnv()
5659
.getConfig()
5760
.getCommonConfig()
61+
.setDegreeOfParallelism(1)
5862
.setMaxTsBlockLineNumber(2)
5963
.setSortBufferSize(128 * 1024)
6064
.setDataRegionGroupExtensionPolicy("CUSTOM")
@@ -110,6 +114,23 @@ public void testUnorderedUnionPreservesEveryDuplicateOnEachNode() throws Excepti
110114
}
111115
}
112116

117+
@Test
118+
public void testPipeVerificationQueryUsesDeterministicPrimaryKeyOrder() throws Exception {
119+
TableModelUtils.createDataBaseAndTable(EnvFactory.getEnv(), "pipe_readings", DATABASE);
120+
TableModelUtils.insertData(DATABASE, "pipe_readings", 0, DEVICES * ROWS, EnvFactory.getEnv());
121+
List<String> expected =
122+
new ArrayList<>(TableModelUtils.generateExpectedResults(0, DEVICES * ROWS));
123+
Collections.sort(expected);
124+
try (Connection connection = EnvFactory.getEnv().getTableConnection();
125+
Statement statement = connection.createStatement()) {
126+
statement.execute("USE " + DATABASE);
127+
try (ResultSet result =
128+
statement.executeQuery(TableModelUtils.getQuerySql("pipe_readings"))) {
129+
TestUtils.assertResultSetEqual(result, TableModelUtils.generateHeaderResults(), expected);
130+
}
131+
}
132+
}
133+
113134
@Test
114135
public void testLegacyTreeAndSharedTreeToTablePaths() throws Exception {
115136
try (Connection connection = EnvFactory.getEnv().getConnection();

‎iotdb-core/calc-commons/src/main/i18n/en/org/apache/iotdb/calc/i18n/CalcMessages.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -582,4 +582,7 @@ private CalcMessages() {}
582582
EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 =
583583
"Aggregate function [%s] requires %d arguments, but got %d";
584584

585+
public static final String
586+
EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 =
587+
"Ready input consumption must be enabled before Collect execution starts";
585588
}

‎iotdb-core/calc-commons/src/main/i18n/zh/org/apache/iotdb/calc/i18n/CalcMessages.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -558,4 +558,7 @@ private CalcMessages() {}
558558
public static final String
559559
EXCEPTION_AGGREGATE_FUNCTION_ARG_REQUIRES_ARG_ARGUMENTS_BUT_GOT_ARG_D6249DD2 =
560560
"聚合函数 [%s] 需要 %d 个参数，但实际传入 %d 个";
561+
public static final String
562+
EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4 =
563+
"必须在 Collect 开始执行前启用就绪输入消费";
561564
}

‎iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/CollectOperator.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import org.apache.iotdb.calc.execution.operator.CommonOperatorContext;
2323
import org.apache.iotdb.calc.execution.operator.Operator;
24+
import org.apache.iotdb.calc.i18n.CalcMessages;
2425
import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper;
2526
import org.apache.iotdb.commons.utils.TestOnly;
2627

@@ -30,6 +31,7 @@
3031

3132
import java.util.List;
3233

34+
import static com.google.common.base.Preconditions.checkState;
3335
import static com.google.common.util.concurrent.Futures.getDone;
3436
import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
3537

@@ -41,6 +43,7 @@ public class CollectOperator implements ProcessOperator {
4143
protected final List<Operator> children;
4244
private AnyChildBlocked anyChildBlocked;
4345
private boolean inited;
46+
private boolean executionStarted;
4447
protected int currentIndex;
4548
private ListenableFuture<Integer> selectionFuture;
4649
private ListenableFuture<Integer> exposedFuture;
@@ -58,6 +61,10 @@ public CollectOperator(CommonOperatorContext operatorContext, List<Operator> chi
5861
*/
5962
public void enableReadyInputConsumption() {
6063
if (anyChildBlocked == null) {
64+
checkState(
65+
!executionStarted,
66+
CalcMessages
67+
.EXCEPTION_READY_INPUT_CONSUMPTION_MUST_BE_ENABLED_BEFORE_COLLECT_EXECUTION_STARTS_5FAF2DC4);
6168
anyChildBlocked = new AnyChildBlocked(children);
6269
}
6370
}
@@ -71,6 +78,7 @@ public boolean hasNext() {
7178

7279
@Override
7380
public TsBlock next() throws Exception {
81+
executionStarted = true;
7482
if (anyChildBlocked == null) {
7583
if (children.get(currentIndex).hasNextWithTimer()) {
7684
TsBlock block = children.get(currentIndex).nextWithTimer();
@@ -137,6 +145,7 @@ protected void closeCurrentChild(int index) throws Exception {
137145

138146
@Override
139147
public ListenableFuture<?> isBlocked() {
148+
executionStarted = true;
140149
if (anyChildBlocked == null) {
141150
if (!inited) {
142151
inited = true;
@@ -170,6 +179,7 @@ public CommonOperatorContext getOperatorContext() {
170179

171180
@Override
172181
public void close() throws Exception {
182+
executionStarted = true;
173183
if (anyChildBlocked == null) {
174184
for (int i = currentIndex; i < children.size(); i++) {
175185
if (children.get(i) != null) {

‎iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import static org.junit.Assert.assertFalse;
5656
import static org.junit.Assert.assertNull;
5757
import static org.junit.Assert.assertSame;
58+
import static org.junit.Assert.assertThrows;
5859
import static org.junit.Assert.assertTrue;
5960
import static org.junit.Assert.fail;
6061

@@ -425,6 +426,119 @@ public void testDefaultMappingCollectRetainsSequentialConsumption() throws Excep
425426
collect.close();
426427
}
427428

429+
@Test
430+
public void testLegacyInitializesOtherInputsOnceAndFinishesAtLastChild() throws Exception {
431+
Child a = new Child(block(1));
432+
Child b = new Child(block(2));
433+
b.ready();
434+
CollectOperator collect =
435+
new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b)));
436+
for (int i = 0; i < 10; i++) {
437+
assertFalse(collect.isBlocked().isDone());
438+
}
439+
assertEquals(11, a.waitCalls);
440+
assertEquals(1, b.waitCalls);
441+
assertEquals(0, b.advances);
442+
a.ready();
443+
assertSame(a.output, collect.next());
444+
a.hasNext = false;
445+
assertNull(collect.next());
446+
assertTrue(collect.hasNext());
447+
assertFalse(collect.isFinished());
448+
for (int i = 0; i < 10; i++) {
449+
assertTrue(collect.isBlocked().isDone());
450+
}
451+
assertEquals(11, a.waitCalls);
452+
assertEquals(11, b.waitCalls);
453+
assertSame(b.output, collect.next());
454+
b.hasNext = false;
455+
assertNull(collect.next());
456+
assertFalse(collect.hasNext());
457+
assertTrue(collect.isFinished());
458+
assertTrue(collect.isBlocked().isDone());
459+
collect.close();
460+
assertEquals(1, a.closes);
461+
assertEquals(1, b.closes);
462+
}
463+
464+
@Test
465+
public void testLegacyMemoryUsesLargestChildPeak() throws Exception {
466+
Child a = new Child(block(1));
467+
Child b = new Child(block(2));
468+
a.peak = 100;
469+
a.retained = 30;
470+
a.returnSize = 10;
471+
b.peak = 80;
472+
b.retained = 50;
473+
b.returnSize = 20;
474+
CollectOperator collect =
475+
new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b)));
476+
assertEquals(100, collect.calculateMaxPeekMemory());
477+
assertEquals(0, collect.calculateRetainedSizeAfterCallingNext());
478+
assertEquals(20, collect.calculateMaxReturnSize());
479+
collect.close();
480+
}
481+
482+
@Test
483+
public void testCannotEnableReadyConsumptionAfterLegacyInputIsRetired() throws Exception {
484+
Child a = new Child(block(1));
485+
Child b = new Child(block(2));
486+
a.hasNext = false;
487+
b.ready();
488+
CollectOperator collect =
489+
new CollectOperator(new TestContext(), new ArrayList<>(Arrays.asList(a, b)));
490+
assertNull(collect.next());
491+
assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption);
492+
assertSame(b.output, collect.next());
493+
b.hasNext = false;
494+
assertNull(collect.next());
495+
assertTrue(collect.isFinished());
496+
}
497+
498+
@Test
499+
public void testCannotEnableReadyConsumptionAfterLegacyReadinessObservation() throws Exception {
500+
Child child = new Child(block(1));
501+
CollectOperator collect =
502+
new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child)));
503+
assertFalse(collect.isBlocked().isDone());
504+
assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption);
505+
collect.close();
506+
}
507+
508+
@Test
509+
public void testCannotEnableReadyConsumptionAfterLegacyNextWithoutIsBlocked() throws Exception {
510+
Child child = new Child(block(1));
511+
child.ready();
512+
CollectOperator collect =
513+
new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child)));
514+
assertSame(child.output, collect.next());
515+
assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption);
516+
collect.close();
517+
}
518+
519+
@Test
520+
public void testCannotEnableReadyConsumptionAfterLegacyClose() throws Exception {
521+
Child child = new Child(block(1));
522+
CollectOperator collect =
523+
new CollectOperator(new TestContext(), new ArrayList<>(Collections.singletonList(child)));
524+
collect.close();
525+
assertThrows(IllegalStateException.class, collect::enableReadyInputConsumption);
526+
assertEquals(1, child.closes);
527+
}
528+
529+
@Test
530+
public void testEnablingReadyConsumptionIsIdempotent() throws Exception {
531+
Child child = new Child(block(1));
532+
child.ready();
533+
CollectOperator collect = collect(child);
534+
collect.enableReadyInputConsumption();
535+
ListenableFuture<?> selection = collect.isBlocked();
536+
collect.enableReadyInputConsumption();
537+
assertSame(selection, collect.isBlocked());
538+
assertSame(child.output, collect.next());
539+
collect.close();
540+
}
541+
428542
@Test
429543
public void testMappingUsesActualChildIdentity() throws Exception {
430544
Child a = new Child(block(10, 11));

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/memory/MemoryPool.java‎

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ private static class FragmentMemoryReservation {
174174
private final AtomicLong inputUsedBytes = new AtomicLong();
175175
private final AtomicLong outputUsedBytes = new AtomicLong();
176176
// Enabled by local planning before drivers start; request limits are supplied by handles.
177-
private boolean sharedBudget;
177+
private volatile boolean sharedBudget;
178178

179179
private AtomicLong getUsedBytes(MemorySide side) {
180180
return side == MemorySide.INPUT ? inputUsedBytes : outputUsedBytes;
@@ -299,6 +299,18 @@ public void deRegisterFragmentInstanceFromQueryMemoryMap(
299299
// forceDeregister == true, after all LocalSourceChannel/LocalSourceHandles are closed.
300300
return;
301301
}
302+
// Pending requests must not outlive their FI registration. Cancel before removing the
303+
// account so another query's free() cannot retry a request against a missing FI. Do not
304+
// hold the registration lock while notifying cancellation listeners.
305+
Iterator<MemoryReservationFuture<Void>> pending = memoryReservationFutures.iterator();
306+
while (pending.hasNext()) {
307+
MemoryReservationFuture<Void> future = pending.next();
308+
if (queryId.equals(future.getQueryId())
309+
&& fragmentInstanceId.equals(future.getFragmentInstanceId())) {
310+
tryCancel(future);
311+
pending.remove();
312+
}
313+
}
302314
synchronized (queryMemoryReservations) {
303315
queryRelatedMemory.remove(fragmentInstanceId);
304316
// 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
451463
DataNodeQueryMessages
452464
.QUERY_EXCEPTION_RELATEDMEMORYRESERVED_CAN_T_BE_NULL_WHEN_FREEING_MEMORY_C80009F2);
453465
}
454-
synchronized (fragment) {
455-
PlanNodeReservation reservation =
456-
requireReservation(
457-
fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId);
466+
PlanNodeReservation reservation =
467+
requireReservation(
468+
fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId);
469+
Object accountingLock = fragment.sharedBudget ? fragment : reservation;
470+
synchronized (accountingLock) {
458471
if (reservation.bytesReserved < bytes) {
459472
throw new IllegalArgumentException(
460473
DataNodeQueryMessages.FREE_MORE_MEMORY_THAN_HAS_BEEN_RESERVED);
@@ -540,10 +553,12 @@ public boolean tryReserve(
540553
long maxBytesCanReserve) {
541554
FragmentMemoryReservation fragment =
542555
getFragmentReservation(queryId, fragmentInstanceId, planNodeId);
543-
synchronized (fragment) {
544-
PlanNodeReservation reservation =
545-
requireReservation(
546-
fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId);
556+
PlanNodeReservation reservation =
557+
requireReservation(
558+
fragment.planNodes.get(planNodeId), queryId, fragmentInstanceId, planNodeId);
559+
// Legacy quotas are independent per plan node. Only shared quotas require the FI lock.
560+
Object accountingLock = fragment.sharedBudget ? fragment : reservation;
561+
synchronized (accountingLock) {
547562
long usedBytes =
548563
fragment.sharedBudget
549564
? fragment.getUsedBytes(reservation.side).get()

0 commit comments

Comments
 (0)