From 0dfce1ce7ff44ac7275500fdcac0160a205400d1 Mon Sep 17 00:00:00 2001 From: JackieTien97 Date: Fri, 9 Oct 2026 16:56:36 +0800 Subject: [PATCH] Allow CollectOperator to consume any ready child --- .../operator/process/AnyChildBlocked.java | 193 ++++++++++ .../operator/process/CollectOperator.java | 79 ++-- .../process/MappingCollectOperator.java | 36 +- .../operator/process/AnyChildBlockedTest.java | 291 +++++++++++++++ .../operator/process/CollectOperatorTest.java | 336 ++++++++++++++++++ 5 files changed, 889 insertions(+), 46 deletions(-) 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/AnyChildBlockedTest.java create mode 100644 iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java 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 0000000000000..19291754db306 --- /dev/null +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlocked.java @@ -0,0 +1,193 @@ +/* + * 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 com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.utils.RamUsageEstimator; + +import java.util.concurrent.ExecutionException; +import java.util.function.IntFunction; + +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + +/** + * Tracks ready children while registering at most one listener per child's pending future. + * + *

Methods are called by the driver thread. Only ReadyListener.run() executes on completion + * threads. The supplier returns null for a closed child. A child's future is refreshed after its + * index is taken, so the caller must perform at most one hasNext/next step before checking + * readiness again. + */ +final class AnyChildBlocked implements AutoCloseable { + private static final long INSTANCE_SIZE = + RamUsageEstimator.shallowSizeOfInstance(AnyChildBlocked.class); + private static final long CHILD_STATE_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ChildState.class); + private static final long LISTENER_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ReadyListener.class); + + private final ChildState[] states; + private final IntFunction> blockedSupplier; + + private int readyIndex = -1; + private SettableFuture waiting; + private boolean closed; + + AnyChildBlocked(int childCount, IntFunction> blockedSupplier) { + this.blockedSupplier = blockedSupplier; + states = new ChildState[childCount]; + for (int i = 0; i < childCount; i++) { + states[i] = new ChildState(i); + } + } + + ListenableFuture isBlocked() { + if (closed || readyIndex >= 0) { + return Futures.immediateVoidFuture(); + } + if (waiting != null && !waiting.isCancelled()) { + return waiting; + } + waiting = null; + + boolean hasChild = false; + for (ChildState state : states) { + if (state.future == null) { + state.future = blockedSupplier.apply(state.index); + } + if (state.future != null) { + hasChild = true; + if (readyIndex < 0 && state.future.isDone()) { + readyIndex = state.index; + } + } + } + + // Initialize all children's futures, even when an earlier child is already ready. + if (readyIndex >= 0 || !hasChild) { + return Futures.immediateVoidFuture(); + } + + SettableFuture round = SettableFuture.create(); + waiting = round; + for (ChildState state : states) { + if (state.future != null) { + state.await(round); + } + } + return round; + } + + /** Returns -1 when no child can be consumed, without blocking the driver thread. */ + int takeReadyChildIndex() throws ExecutionException { + if (!isBlocked().isDone() || closed) { + return -1; + } + int index = readyIndex; + if (index < 0 && waiting != null) { + index = Futures.getDone(waiting); + } + if (index < 0) { + return -1; + } + + // Completion may also mean failure/cancellation; propagate it before accessing the child. + Futures.getDone(states[index].future); + states[index].reset(); + readyIndex = -1; + waiting = null; + return index; + } + + @Override + public void close() { + closed = true; + readyIndex = -1; + for (ChildState state : states) { + state.reset(); + } + if (waiting != null) { + waiting.set(-1); + waiting = null; + } + } + + long ramBytesUsed() { + long size = INSTANCE_SIZE + RamUsageEstimator.shallowSizeOf(states); + for (ChildState state : states) { + size += CHILD_STATE_SIZE; + if (state.listener != null) { + size += LISTENER_SIZE; + } + } + return size; + } + + private static final class ChildState { + private final int index; + private ListenableFuture future; + private ReadyListener listener; + + private ChildState(int index) { + this.index = index; + } + + private void await(SettableFuture round) { + if (listener == null) { + listener = new ReadyListener(index, round); + future.addListener(listener, directExecutor()); + } else { + listener.target = round; + } + // Covers completion before/during retargeting, even if the old callback has already run. + if (future.isDone()) { + round.set(index); + } + } + + private void reset() { + future = null; + if (listener != null) { + listener.target = null; + listener = null; + } + } + } + + private static final class ReadyListener implements Runnable { + private final int index; + private volatile SettableFuture target; + + private ReadyListener(int index, SettableFuture target) { + this.index = index; + this.target = target; + } + + @Override + public void run() { + SettableFuture round = target; + if (round != null) { + round.set(index); + } + } + } +} 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 921d273d04e8a..6d4a66532cf0b 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 @@ -36,7 +36,8 @@ public class CollectOperator implements ProcessOperator { private final CommonOperatorContext operatorContext; protected final List children; - private boolean inited = false; + private final AnyChildBlocked childBlocker; + private int remainingChildren; protected int currentIndex; @@ -44,46 +45,55 @@ public CollectOperator(CommonOperatorContext operatorContext, List chi this.operatorContext = operatorContext; this.children = children; this.currentIndex = 0; + this.remainingChildren = children.size(); + this.childBlocker = + new AnyChildBlocked( + children.size(), + index -> { + Operator child = this.children.get(index); + return child == null ? null : child.isBlocked(); + }); } @Override public boolean hasNext() throws Exception { - return currentIndex < children.size(); + return remainingChildren > 0; } @Override public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - return children.get(currentIndex).nextWithTimer(); - } else { - closeCurrentChild(currentIndex); - currentIndex++; + int index = childBlocker.takeReadyChildIndex(); + if (index < 0) { return null; } + + // Keep the selected index on the driver thread, including while subclasses map the result. + currentIndex = index; + Operator child = children.get(index); + if (child.hasNextWithTimer()) { + return child.nextWithTimer(); + } + closeCurrentChild(index); + return null; } protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); - children.set(index, null); + Operator child = children.get(index); + if (child != null) { + child.close(); + children.set(index, null); + remainingChildren--; + } } @Override public ListenableFuture isBlocked() { - if (!inited) { - inited = true; - for (Operator child : children) { - child.isBlocked(); - } - } - if (currentIndex >= children.size()) { - return NOT_BLOCKED; - } - return children.get(currentIndex).isBlocked(); + return childBlocker.isBlocked(); } @Override public boolean isFinished() throws Exception { - return currentIndex >= children.size(); + return remainingChildren == 0; } @Override @@ -93,7 +103,8 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { - for (int i = currentIndex, n = children.size(); i < n; i++) { + childBlocker.close(); + for (int i = 0, n = children.size(); i < n; i++) { Operator currentChild = children.get(i); if (currentChild != null) { closeCurrentChild(i); @@ -103,9 +114,18 @@ public void close() throws Exception { @Override public long calculateMaxPeekMemory() { - long maxPeekMemory = 0; + // Switching away from a blocked child may leave its intermediate state retained. + long retainedSize = calculateRetainedSizeAfterCallingNext(); + long maxPeekMemory = retainedSize; for (Operator child : children) { - maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + if (child != null) { + maxPeekMemory = + Math.max( + maxPeekMemory, + retainedSize + - child.calculateRetainedSizeAfterCallingNext() + + child.calculateMaxPeekMemoryWithCounter()); + } } return maxPeekMemory; } @@ -114,14 +134,22 @@ public long calculateMaxPeekMemory() { 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 retainedSize = 0; + for (Operator child : children) { + if (child != null) { + retainedSize += child.calculateRetainedSizeAfterCallingNext(); + } + } + return retainedSize; } @TestOnly @@ -132,6 +160,7 @@ public List getChildren() { @Override public long ramBytesUsed() { return INSTANCE_SIZE + + childBlocker.ramBytesUsed() + children.stream() .mapToLong(MemoryEstimationHelper::getEstimatedSizeOfAccountableObject) .sum() 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 8bc1c18e9e66a..51cd27e36bdb2 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 @@ -50,32 +50,26 @@ public MappingCollectOperator( @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++; + TsBlock tsBlock = super.next(); + if (tsBlock == null) { return null; } + + 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); } + @Override protected void closeCurrentChild(int index) throws Exception { - children.get(index).close(); - children.set(index, null); + super.closeCurrentChild(index); mappings.set(index, null); } } diff --git a/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlockedTest.java b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlockedTest.java new file mode 100644 index 0000000000000..9cd0680c95ac2 --- /dev/null +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/AnyChildBlockedTest.java @@ -0,0 +1,291 @@ +/* + * 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 com.google.common.util.concurrent.ForwardingListenableFuture; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; +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 static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static org.junit.Assert.assertTrue; + +public class AnyChildBlockedTest { + @Test(timeout = 30000) + public void testFastPath() throws Exception { + CountingFuture ready = new CountingFuture(); + ready.complete(); + CountingFuture slow = new CountingFuture(); + CountingFuture[] inputs = {ready, slow}; + int[] polls = {0, 0}; + AnyChildBlocked blocker = + new AnyChildBlocked( + 2, + i -> { + polls[i]++; + return inputs[i]; + }); + for (int i = 0; i < 100; i++) { + check(blocker.isBlocked().isDone(), "Ready input should avoid waiting"); + } + check(polls[0] == 1 && polls[1] == 1, "Initialize every input and reuse cached futures"); + check(ready.registrations == 0 && slow.registrations == 0, "Fast path registers no listeners"); + check(blocker.takeReadyChildIndex() == 0, "Ready index"); + blocker.close(); + } + + @Test(timeout = 30000) + public void testRepeatedWaiting() throws Exception { + CountingFuture first = new CountingFuture(); + CountingFuture second = new CountingFuture(); + CountingFuture[] inputs = {first, second}; + AnyChildBlocked blocker = new AnyChildBlocked(2, i -> inputs[i]); + ListenableFuture round = blocker.isBlocked(); + for (int i = 0; i < 100; i++) { + check(blocker.isBlocked() == round, "Reuse pending aggregate"); + } + check(first.registrations == 1 && second.registrations == 1, "Register once per child"); + second.complete(); + first.complete(); + check(blocker.takeReadyChildIndex() == 1, "Later completion must not overwrite winner"); + inputs[1] = null; + check(blocker.takeReadyChildIndex() == 0, "Do not lose another completed child"); + blocker.close(); + } + + @Test(timeout = 30000) + public void testLongBlockedChild() throws Exception { + CountingFuture slow = new CountingFuture(); + CountingFuture[] inputs = {slow, new CountingFuture()}; + int[] polls = {0, 0}; + AnyChildBlocked blocker = + new AnyChildBlocked( + 2, + i -> { + polls[i]++; + return inputs[i]; + }); + for (int i = 0; i < 2000; i++) { + ListenableFuture round = blocker.isBlocked(); + check(!round.isDone(), "All inputs block"); + inputs[1].complete(); + check(round.isDone(), "One child wakes the collector"); + check(blocker.takeReadyChildIndex() == 1, "Identify the fast child"); + inputs[1] = new CountingFuture(); + } + check( + slow.registrations == 1 && polls[0] == 1, + "Slow child has one future and listener across rounds"); + ListenableFuture last = blocker.isBlocked(); + slow.complete(); + check( + last.isDone() && blocker.takeReadyChildIndex() == 0, "Slow child wakes the current round"); + blocker.close(); + } + + @Test(timeout = 30000) + public void testRearmRaceAndStaleCallback() throws Exception { + CountingFuture slow = new CountingFuture(); + slow.defer = true; + CountingFuture[] inputs = {slow, new CountingFuture()}; + AnyChildBlocked blocker = new AnyChildBlocked(2, i -> inputs[i]); + blocker.isBlocked(); + inputs[1].complete(); + check(blocker.takeReadyChildIndex() == 1, "Consume other child first"); + inputs[1] = new CountingFuture(); + slow.afterProbe = slow::complete; + ListenableFuture next = blocker.isBlocked(); + check(next.isDone(), "Recheck after retargeting must recover a completion after the scan"); + check(blocker.takeReadyChildIndex() == 0, "Recover the correct child"); + inputs[0] = new CountingFuture(); + ListenableFuture following = blocker.isBlocked(); + slow.flushCallbacks(); + check(!following.isDone(), "Delayed old callback must not wake a new child generation"); + inputs[1].complete(); + check( + following.isDone() && blocker.takeReadyChildIndex() == 1, "Current generation still wakes"); + check(slow.registrations == 1, "Rearming adds no listener"); + blocker.close(); + } + + @Test(timeout = 30000) + public void testCompletionDuringRegistration() throws Exception { + CountingFuture input = new CountingFuture(); + input.beforeRegister = input::complete; + AnyChildBlocked blocker = new AnyChildBlocked(1, i -> input); + check(blocker.isBlocked().isDone(), "Completion between scan and addListener must not be lost"); + check(blocker.takeReadyChildIndex() == 0, "Identify raced completion"); + blocker.close(); + } + + @Test(timeout = 30000) + public void testFailureAndCancellation() throws Exception { + CountingFuture input = new CountingFuture(); + AnyChildBlocked blocker = new AnyChildBlocked(1, i -> input); + ListenableFuture abandoned = blocker.isBlocked(); + abandoned.cancel(false); + check(!input.isCancelled(), "Canceling a wait must not cancel the child"); + ListenableFuture replacement = blocker.isBlocked(); + check( + replacement != abandoned && input.registrations == 1, + "Canceled waits rearm existing listeners"); + IllegalArgumentException cause = new IllegalArgumentException(); + input.delegate.setException(cause); + check(replacement.isDone(), "Failure wakes the collector"); + try { + blocker.takeReadyChildIndex(); + throw new AssertionError("Expected child failure"); + } catch (ExecutionException expected) { + check(expected.getCause() == cause, "Preserve original cause"); + } + blocker.close(); + CountingFuture canceled = new CountingFuture(); + AnyChildBlocked second = new AnyChildBlocked(1, i -> canceled); + second.isBlocked(); + canceled.cancel(false); + try { + second.takeReadyChildIndex(); + throw new AssertionError("Expected child cancellation"); + } catch (CancellationException expected) { + // Expected. + } + second.close(); + } + + @Test(timeout = 30000) + public void testEmptyAndClosed() throws Exception { + AnyChildBlocked empty = new AnyChildBlocked(0, i -> null); + check( + empty.isBlocked().isDone() && empty.takeReadyChildIndex() == -1, "Empty input is finished"); + CountingFuture input = new CountingFuture(); + AnyChildBlocked blocker = new AnyChildBlocked(2, i -> i == 0 ? null : input); + ListenableFuture pending = blocker.isBlocked(); + check(blocker.takeReadyChildIndex() == -1, "Never block the driver in takeReadyChildIndex"); + blocker.close(); + check(pending.isDone() && blocker.takeReadyChildIndex() == -1, "Close releases waiters"); + check(!input.isCancelled(), "Close does not cancel child futures"); + input.complete(); + check(blocker.takeReadyChildIndex() == -1, "Completion after close is harmless"); + } + + @Test(timeout = 30000) + public void testConcurrentCompletion() throws Exception { + ExecutorService workers = Executors.newFixedThreadPool(2); + try { + for (int round = 0; round < 1000; round++) { + CountingFuture[] inputs = {new CountingFuture(), new CountingFuture()}; + CountingFuture firstInput = inputs[0]; + CountingFuture secondInput = inputs[1]; + AnyChildBlocked blocker = new AnyChildBlocked(2, i -> inputs[i]); + CountDownLatch start = new CountDownLatch(1); + Future first = workers.submit(() -> completeAfter(start, firstInput)); + Future second = workers.submit(() -> completeAfter(start, secondInput)); + start.countDown(); + blocker.isBlocked().get(5, TimeUnit.SECONDS); + int winner = blocker.takeReadyChildIndex(); + check(winner == 0 || winner == 1, "Concurrent completion has a valid winner"); + inputs[winner] = null; + first.get(5, TimeUnit.SECONDS); + second.get(5, TimeUnit.SECONDS); + check( + blocker.takeReadyChildIndex() == 1 - winner, "Retain the other concurrent completion"); + blocker.close(); + } + } finally { + workers.shutdownNow(); + } + } + + private static void completeAfter(CountDownLatch start, CountingFuture input) { + try { + start.await(); + input.complete(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + private static void check(boolean condition, String message) { + assertTrue(message, condition); + } + + private static final class CountingFuture extends ForwardingListenableFuture { + private final SettableFuture delegate = SettableFuture.create(); + private int registrations; + private boolean defer; + private final List deferred = new ArrayList<>(); + private Runnable afterProbe; + private Runnable beforeRegister; + + @Override + protected ListenableFuture delegate() { + return delegate; + } + + @Override + public boolean isDone() { + boolean done = delegate.isDone(); + Runnable hook = afterProbe; + afterProbe = null; + if (hook != null) { + hook.run(); + } + return done; + } + + @Override + public void addListener(Runnable listener, Executor executor) { + registrations++; + if (beforeRegister != null) { + beforeRegister.run(); + } + delegate.addListener( + () -> { + if (defer) { + deferred.add(() -> executor.execute(listener)); + } else { + executor.execute(listener); + } + }, + directExecutor()); + } + + private void complete() { + delegate.set(null); + } + + private void flushCallbacks() { + deferred.forEach(Runnable::run); + deferred.clear(); + } + } +} 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 0000000000000..555103033ffeb --- /dev/null +++ b/iotdb-core/calc-commons/src/test/java/org/apache/iotdb/calc/execution/operator/process/CollectOperatorTest.java @@ -0,0 +1,336 @@ +/* + * 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 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.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Deque; +import java.util.List; + +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; + +public class CollectOperatorTest { + @Test + public void testConsumesReadyChildAndFinishesOutOfOrder() throws Exception { + TestingOperator first = new TestingOperator(); + TestingOperator second = new TestingOperator(); + CollectOperator collector = collect(first, second); + ListenableFuture waiting = collector.isBlocked(); + assertFalse(waiting.isDone()); + + TsBlock secondBlock = block(20); + second.addBlock(secondBlock); + assertTrue(waiting.isDone()); + assertSame(secondBlock, collector.next()); + assertEquals(1, collector.currentIndex); + + second.finish(); + assertNull(collector.next()); + assertNull(collector.getChildren().get(1)); + assertEquals(1, second.closeCalls); + assertTrue(collector.hasNext()); + assertFalse(collector.isFinished()); + assertFalse(collector.isBlocked().isDone()); + + TsBlock firstBlock = block(10); + first.addBlock(firstBlock); + assertSame(firstBlock, collector.next()); + first.finish(); + assertNull(collector.next()); + assertFalse(collector.hasNext()); + assertTrue(collector.isFinished()); + assertTrue(collector.isBlocked().isDone()); + assertNull(collector.next()); + collector.close(); + assertEquals(1, first.closeCalls); + assertEquals(1, second.closeCalls); + } + + @Test + public void testCloseIncludesEarlierChildrenAndReleasesWaiting() throws Exception { + TestingOperator first = new TestingOperator(); + TestingOperator second = new TestingOperator(); + CollectOperator collector = collect(first, second); + TsBlock result = block(20); + second.addBlock(result); + assertSame(result, collector.next()); + ListenableFuture waiting = collector.isBlocked(); + assertFalse(waiting.isDone()); + + collector.close(); + assertTrue(waiting.isDone()); + assertTrue(collector.isFinished()); + assertFalse(collector.hasNext()); + first.finish(); + second.finish(); + assertNull(collector.next()); + collector.close(); + assertEquals(1, first.closeCalls); + assertEquals(1, second.closeCalls); + } + + @Test + public void testYieldRefreshesOnlyConsumedChild() throws Exception { + TestingOperator first = new TestingOperator(); + TestingOperator second = new TestingOperator(); + CollectOperator collector = collect(first, second); + collector.isBlocked(); + second.blocked.set(null); + assertNull(collector.next()); + ListenableFuture waiting = collector.isBlocked(); + assertFalse(waiting.isDone()); + assertEquals(1, first.blockedCalls); + assertEquals(2, second.blockedCalls); + assertSame(waiting, collector.isBlocked()); + + TsBlock result = block(20); + second.addBlock(result); + assertSame(result, collector.next()); + assertEquals(0, first.nextCalls); + assertEquals(2, second.nextCalls); + collector.close(); + } + + @Test + public void testMappingUsesReadyChildAndClosesOutOfOrder() throws Exception { + TestingOperator first = new TestingOperator(); + TestingOperator second = new TestingOperator(); + List> mappings = + new ArrayList<>(Arrays.asList(Arrays.asList(0, 1), Arrays.asList(1, 0))); + MappingCollectOperator collector = + new MappingCollectOperator(null, children(first, second), mappings); + ListenableFuture waiting = collector.isBlocked(); + second.addBlock(block(20, 21)); + assertTrue(waiting.isDone()); + TsBlock result = collector.next(); + assertEquals(21, result.getColumn(0).getInt(0)); + assertEquals(20, result.getColumn(1).getInt(0)); + + second.finish(); + assertNull(collector.next()); + assertNull(mappings.get(1)); + assertTrue(collector.hasNext()); + first.addBlock(block(10, 11)); + result = collector.next(); + assertEquals(10, result.getColumn(0).getInt(0)); + assertEquals(11, result.getColumn(1).getInt(0)); + first.finish(); + assertNull(collector.next()); + assertNull(mappings.get(0)); + assertTrue(collector.isFinished()); + collector.close(); + assertEquals(1, first.closeCalls); + assertEquals(1, second.closeCalls); + } + + @Test + public void testAnotherCompletionCannotChangeSelectedMapping() throws Exception { + TestingOperator first = new TestingOperator(); + TestingOperator second = new TestingOperator(); + MappingCollectOperator collector = + new MappingCollectOperator( + null, + children(first, second), + new ArrayList<>(Arrays.asList(Arrays.asList(0, 1), Arrays.asList(1, 0)))); + collector.isBlocked(); + second.addBlock(block(20, 21)); + second.onNext = () -> first.addBlock(block(10, 11)); + TsBlock result = collector.next(); + assertEquals(21, result.getColumn(0).getInt(0)); + assertEquals(20, result.getColumn(1).getInt(0)); + assertEquals(1, collector.currentIndex); + result = collector.next(); + assertEquals(10, result.getColumn(0).getInt(0)); + assertEquals(11, result.getColumn(1).getInt(0)); + collector.close(); + } + + @Test + public void testMemoryIncludesOtherRetainedChildren() throws Exception { + TestingOperator first = new TestingOperator(); + first.peekMemory = 100; + first.retainedMemory = 30; + first.returnSize = 40; + TestingOperator second = new TestingOperator(); + second.peekMemory = 200; + second.retainedMemory = 50; + second.returnSize = 60; + CollectOperator collector = collect(first, second); + assertEquals(80, collector.calculateRetainedSizeAfterCallingNext()); + assertEquals(230, collector.calculateMaxPeekMemory()); + assertEquals(60, collector.calculateMaxReturnSize()); + assertTrue(collector.ramBytesUsed() > 0); + + first.finish(); + assertNull(collector.next()); + assertEquals(50, collector.calculateRetainedSizeAfterCallingNext()); + assertEquals(200, collector.calculateMaxPeekMemory()); + assertEquals(60, collector.calculateMaxReturnSize()); + collector.close(); + assertEquals(0, collector.calculateRetainedSizeAfterCallingNext()); + assertEquals(0, collector.calculateMaxPeekMemory()); + assertEquals(0, collector.calculateMaxReturnSize()); + } + + @Test + public void testEmptyChildren() throws Exception { + CollectOperator collector = collect(); + assertTrue(collector.isFinished()); + assertFalse(collector.hasNext()); + assertTrue(collector.isBlocked().isDone()); + assertNull(collector.next()); + collector.close(); + } + + private static CollectOperator collect(TestingOperator... children) { + return new CollectOperator(null, children(children)); + } + + private static List children(TestingOperator... children) { + return new ArrayList<>(Arrays.asList(children)); + } + + 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 final class TestingOperator implements Operator { + private SettableFuture blocked = SettableFuture.create(); + private final Deque blocks = new ArrayDeque<>(); + private boolean finished; + private int blockedCalls; + private int nextCalls; + private int closeCalls; + private long peekMemory; + private long retainedMemory; + private long returnSize; + private Runnable onNext; + + private void addBlock(TsBlock block) { + blocks.addLast(block); + blocked.set(null); + } + + private void finish() { + finished = true; + blocked.set(null); + } + + @Override + public CommonOperatorContext getOperatorContext() { + return null; + } + + @Override + public ListenableFuture isBlocked() { + blockedCalls++; + return blocked; + } + + @Override + public boolean hasNext() { + assertTrue("Do not probe blocked children", blocked.isDone()); + return !finished || !blocks.isEmpty(); + } + + @Override + public boolean hasNextWithTimer() { + return hasNext(); + } + + @Override + public TsBlock next() { + assertTrue("Do not consume blocked children", blocked.isDone()); + nextCalls++; + TsBlock block = blocks.pollFirst(); + blocked = SettableFuture.create(); + if (finished || !blocks.isEmpty()) { + blocked.set(null); + } + if (onNext != null) { + onNext.run(); + } + return block; + } + + @Override + public TsBlock nextWithTimer() { + return next(); + } + + @Override + public void close() { + closeCalls++; + } + + @Override + public boolean isFinished() { + return finished && blocks.isEmpty(); + } + + @Override + public long calculateMaxPeekMemory() { + return peekMemory; + } + + @Override + public long calculateMaxPeekMemoryWithCounter() { + return peekMemory; + } + + @Override + public long calculateMaxReturnSize() { + return returnSize; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return retainedMemory; + } + + @Override + public long ramBytesUsed() { + return 64; + } + } +}