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..133a8cda3c818 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 @@ -25,11 +25,14 @@ import org.apache.iotdb.commons.utils.TestOnly; import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.utils.RamUsageEstimator; import java.util.List; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + public class CollectOperator implements ProcessOperator { private static final long INSTANCE_SIZE = RamUsageEstimator.shallowSizeOfInstance(CollectOperator.class); @@ -38,28 +41,44 @@ public class CollectOperator implements ProcessOperator { protected final List children; private boolean inited = false; - protected int currentIndex; + /** The next child considered by the fair ready-first scheduler. */ + protected int nextCursor; public CollectOperator(CommonOperatorContext operatorContext, List children) { this.operatorContext = operatorContext; this.children = children; - this.currentIndex = 0; + this.nextCursor = 0; } @Override public boolean hasNext() throws Exception { - return currentIndex < children.size(); + for (Operator child : children) { + if (child != null && child.hasNextWithTimer()) { + return true; + } + } + return false; } @Override public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - return children.get(currentIndex).nextWithTimer(); - } else { - closeCurrentChild(currentIndex); - currentIndex++; - return null; + for (int offset = 0; offset < children.size(); offset++) { + int index = (nextCursor + offset) % children.size(); + Operator child = children.get(index); + if (child == null) { + continue; + } + if (!child.hasNextWithTimer() || child.isFinished()) { + closeCurrentChild(index); + continue; + } + if (!child.isBlocked().isDone()) { + continue; + } + nextCursor = (index + 1) % children.size(); + return child.nextWithTimer(); } + return null; } protected void closeCurrentChild(int index) throws Exception { @@ -72,18 +91,38 @@ public ListenableFuture isBlocked() { if (!inited) { inited = true; for (Operator child : children) { - child.isBlocked(); + if (child != null) { + child.isBlocked(); + } } } - if (currentIndex >= children.size()) { + if (children.isEmpty()) { return NOT_BLOCKED; } - return children.get(currentIndex).isBlocked(); + SettableFuture blocked = SettableFuture.create(); + boolean hasBlockedChild = false; + for (Operator child : children) { + if (child == null) { + continue; + } + ListenableFuture childBlocked = child.isBlocked(); + if (childBlocked.isDone()) { + return NOT_BLOCKED; + } + hasBlockedChild = true; + childBlocked.addListener(() -> blocked.set(null), directExecutor()); + } + return hasBlockedChild ? blocked : NOT_BLOCKED; } @Override public boolean isFinished() throws Exception { - return currentIndex >= children.size(); + for (Operator child : children) { + if (child != null && !child.isFinished()) { + return false; + } + } + return true; } @Override @@ -93,7 +132,7 @@ public CommonOperatorContext getOperatorContext() { @Override public void close() throws Exception { - for (int i = currentIndex, n = children.size(); i < n; i++) { + for (int i = 0, n = children.size(); i < n; i++) { Operator currentChild = children.get(i); if (currentChild != null) { closeCurrentChild(i); @@ -103,25 +142,32 @@ public void close() throws Exception { @Override public long calculateMaxPeekMemory() { - long maxPeekMemory = 0; + long totalPeekMemory = 0; for (Operator child : children) { - maxPeekMemory = Math.max(maxPeekMemory, child.calculateMaxPeekMemoryWithCounter()); + if (child != null) { + totalPeekMemory += child.calculateMaxPeekMemoryWithCounter(); + } } - return maxPeekMemory; + return totalPeekMemory; } @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; + return children.stream() + .filter(child -> child != null) + .mapToLong(Operator::calculateRetainedSizeAfterCallingNext) + .sum(); } @TestOnly 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..2e6dfea134e3d 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,27 +50,35 @@ public MappingCollectOperator( @Override public TsBlock next() throws Exception { - if (children.get(currentIndex).hasNextWithTimer()) { - TsBlock tsBlock = children.get(currentIndex).nextWithTimer(); + for (int offset = 0; offset < children.size(); offset++) { + int index = (nextCursor + offset) % children.size(); + if (children.get(index) == null) { + continue; + } + if (!children.get(index).hasNextWithTimer() || children.get(index).isFinished()) { + closeCurrentChild(index); + continue; + } + if (!children.get(index).isBlocked().isDone()) { + continue; + } + nextCursor = (index + 1) % children.size(); + TsBlock tsBlock = children.get(index).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; + Column[] columns = new Column[outputColumnsCount]; + List mapping = mappings.get(index); + 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); } + return null; } protected void closeCurrentChild(int index) throws Exception { diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/HashAggregationOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/HashAggregationOperator.java index 7a58000f5ab88..36e0371e6bad1 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/HashAggregationOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/HashAggregationOperator.java @@ -178,11 +178,13 @@ private void closeAggregationBuilder() { // outputPages = null; if (aggregationBuilder != null) { aggregationBuilder.close(); - // aggregationBuilder.close() will release all memory reserved in memory accounting. - // The reference must be set to null afterwards to avoid unaccounted memory. aggregationBuilder = null; } - // memoryContext.setBytes(0); + if (previousRetainedSize > 0) { + memoryReservationManager.releaseMemoryCumulatively(previousRetainedSize); + memoryReservationManager.releaseMemoryImmediately(); + previousRetainedSize = 0; + } } @Override @@ -192,6 +194,7 @@ public boolean isFinished() throws Exception { @Override public void close() throws Exception { + closeAggregationBuilder(); child.close(); aggregators.forEach(GroupedAggregator::close); } diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/StreamingHashAggregationOperator.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/StreamingHashAggregationOperator.java index b85a2d509737b..b267055b41dda 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/StreamingHashAggregationOperator.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/source/relational/aggregation/grouped/StreamingHashAggregationOperator.java @@ -306,11 +306,13 @@ private void closeAggregationBuilder() { // outputPages = null; if (aggregationBuilder != null) { aggregationBuilder.close(); - // aggregationBuilder.close() will release all memory reserved in memory accounting. - // The reference must be set to null afterward to avoid unaccounted memory. aggregationBuilder = null; } - // memoryContext.setBytes(0); + if (previousRetainedSize > 0) { + memoryReservationManager.releaseMemoryCumulatively(previousRetainedSize); + memoryReservationManager.releaseMemoryImmediately(); + previousRetainedSize = 0; + } } @Override @@ -320,6 +322,7 @@ public boolean isFinished() throws Exception { @Override public void close() throws Exception { + closeAggregationBuilder(); child.close(); } diff --git a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/memory/MemoryReservationManager.java b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/memory/MemoryReservationManager.java index 18ad895bb3146..84fb14b8833da 100644 --- a/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/memory/MemoryReservationManager.java +++ b/iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/planner/memory/MemoryReservationManager.java @@ -54,6 +54,11 @@ public interface MemoryReservationManager { */ void releaseMemoryImmediately(final long size); + /** Flushes any batched release accumulated by {@link #releaseMemoryCumulatively(long)}. */ + default void releaseMemoryImmediately() { + // Implementations that batch releases override this method. + } + /** * Release all reserved memory immediately. Make sure this method is called when the lifecycle of * this manager ends, Or the memory to be released in the batch may not be released correctly. 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 5812d4e37cdd3..e33ad3ea43d86 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 @@ -57,6 +57,7 @@ public class SharedTsBlockQueue { private final TFragmentInstanceId localFragmentInstanceId; private final String localPlanNodeId; + private String memoryReservationPlanNodeId; private final String fullFragmentInstanceId; @@ -78,6 +79,11 @@ public class SharedTsBlockQueue { private ListenableFuture blockedOnMemory; + /** Limits a grouped local source to one in-flight block until the consumer removes it. */ + private SettableFuture sharedQuotaBlocked; + + private boolean usesSharedMemoryReservationGroup; + private volatile boolean closed = false; private boolean alreadyRegistered = false; @@ -117,6 +123,7 @@ public SharedTsBlockQueue( this.localPlanNodeId = Validate.notNull( planNodeId, DataNodeQueryMessages.EXCEPTION_PLANNODE_ID_CANNOT_BE_NULL_F91303CD); + this.memoryReservationPlanNodeId = this.localPlanNodeId; this.localMemoryManager = Validate.notNull( localMemoryManager, @@ -144,6 +151,18 @@ public void setMaxBytesCanReserve(long maxBytesCanReserve) { this.maxBytesCanReserve = maxBytesCanReserve; } + public void setMemoryReservationGroupId(String reservationPlanNodeId) { + this.memoryReservationPlanNodeId = Validate.notNull(reservationPlanNodeId); + this.usesSharedMemoryReservationGroup = !localPlanNodeId.equals(reservationPlanNodeId); + if (usesSharedMemoryReservationGroup) { + // The planner may first apply the legacy per-exchange limit before it knows this queue's + // Collect group. Restore the fragment-instance limit so the shared reservation key, rather + // than an individual queue, is the group-wide limit. + this.maxBytesCanReserve = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + } + } + public long getMaxBytesCanReserve() { return maxBytesCanReserve; } @@ -228,7 +247,7 @@ public TsBlock remove() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, reservedBytes); bufferRetainedSizeInBytes -= reservedBytes; } @@ -238,6 +257,11 @@ public TsBlock remove() { if (sinkChannel != null) { sinkChannel.checkAndInvokeOnFinished(); } + if (usesSharedMemoryReservationGroup + && sharedQuotaBlocked != null + && !sharedQuotaBlocked.isDone()) { + sharedQuotaBlocked.set(null); + } if (blocked.isDone() && queue.isEmpty() && !noMoreTsBlocks) { blocked = SettableFuture.create(); } @@ -261,11 +285,16 @@ public ListenableFuture add(TsBlock tsBlock) { Validate.isTrue( blockedOnMemory == null || blockedOnMemory.isDone(), DataNodeQueryMessages.EXCEPTION_SHAREDTSBLOCKQUEUE_IS_FULL_87493E26); + Validate.isTrue( + !usesSharedMemoryReservationGroup + || sharedQuotaBlocked == null + || sharedQuotaBlocked.isDone(), + DataNodeQueryMessages.EXCEPTION_SHAREDTSBLOCKQUEUE_IS_FULL_87493E26); if (!alreadyRegistered) { localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, fullFragmentInstanceId, memoryReservationPlanNodeId); alreadyRegistered = true; } MemoryReservationResult reserveResult = @@ -274,8 +303,8 @@ public ListenableFuture add(TsBlock tsBlock) { .reserveWithPriority( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, - tsBlock.getSizeInBytes(), + memoryReservationPlanNodeId, + tsBlock.getRetainedSizeInBytes(), maxBytesCanReserve, isHighestPriority); blockedOnMemory = reserveResult.getFuture(); @@ -302,7 +331,11 @@ public ListenableFuture add(TsBlock tsBlock) { if (!blocked.isDone()) { blocked.set(null); } - channelBlocked.set(null); + if (usesSharedMemoryReservationGroup) { + sharedQuotaBlocked = channelBlocked; + } else { + channelBlocked.set(null); + } } }, // Use directExecutor() here could lead to deadlock. Thread A holds lock of @@ -319,6 +352,10 @@ public ListenableFuture add(TsBlock tsBlock) { if (!blocked.isDone()) { blocked.set(null); } + if (usesSharedMemoryReservationGroup) { + sharedQuotaBlocked = SettableFuture.create(); + return sharedQuotaBlocked; + } return blockedOnMemory; } } @@ -338,6 +375,9 @@ public void close() { if (blockedOnMemory != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } + if (sharedQuotaBlocked != null && !sharedQuotaBlocked.isDone()) { + sharedQuotaBlocked.set(null); + } queue.clear(); if (bufferRetainedSizeInBytes > 0L) { localMemoryManager @@ -345,7 +385,7 @@ public void close() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } @@ -382,6 +422,9 @@ public void abort() { if (blockedOnMemory != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } + if (sharedQuotaBlocked != null && !sharedQuotaBlocked.isDone()) { + sharedQuotaBlocked.set(null); + } queue.clear(); if (bufferRetainedSizeInBytes > 0L) { localMemoryManager @@ -389,7 +432,7 @@ public void abort() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } @@ -411,6 +454,9 @@ public void abort(Throwable t) { if (blockedOnMemory != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blockedOnMemory); } + if (sharedQuotaBlocked != null && !sharedQuotaBlocked.isDone()) { + sharedQuotaBlocked.set(null); + } queue.clear(); if (bufferRetainedSizeInBytes > 0L) { localMemoryManager @@ -418,7 +464,7 @@ public void abort(Throwable t) { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java index 32e5557b8f922..7ed7e342abd56 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/ISourceHandle.java @@ -90,4 +90,9 @@ public interface ISourceHandle extends Accountable { /** Set max bytes this handle can reserve from memory pool. */ void setMaxBytesCanReserve(long maxBytesCanReserve); + + /** Use a shared memory reservation key with the sibling inputs of a Collect node. */ + default void setMemoryReservationGroupId(String reservationPlanNodeId) { + // Handles that do not reserve query memory can ignore the group. + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java index 544421712669a..83713b40e1705 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java @@ -306,8 +306,14 @@ public SharedTsBlockQueue getSharedTsBlockQueue() { @Override public void setMaxBytesCanReserve(long maxBytesCanReserve) { - // do nothing, the maxBytesCanReserve of SharedTsBlockQueue should be set by corresponding - // LocalSinkChannel + // The queue limit is configured by its corresponding LocalSinkChannel. A grouped queue + // restores its fragment-instance limit in setMemoryReservationGroupId before the shared + // reservation key is applied. + } + + @Override + public void setMemoryReservationGroupId(String reservationPlanNodeId) { + queue.setMemoryReservationGroupId(reservationPlanNodeId); } @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 aed3f9533fb1d..732c70d85a065 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 @@ -85,6 +85,8 @@ public class SourceHandle implements ISourceHandle { private final String fullFragmentInstanceId; private final String localPlanNodeId; + private String memoryReservationPlanNodeId; + private boolean usesSharedMemoryReservationGroup; private final int indexOfUpstreamSinkHandle; private final LocalMemoryManager localMemoryManager; @@ -200,6 +202,8 @@ public SourceHandle( Validate.notNull( localPlanNodeId, DataNodeQueryMessages.EXCEPTION_LOCALPLANNODEID_CAN_NOT_BE_NULL_DOT_44A34A33); + this.memoryReservationPlanNodeId = this.localPlanNodeId; + this.usesSharedMemoryReservationGroup = false; this.indexOfUpstreamSinkHandle = indexOfUpstreamSinkHandle; this.localMemoryManager = Validate.notNull( @@ -223,7 +227,7 @@ public SourceHandle( localMemoryManager .getQueryPool() .registerPlanNodeIdToQueryMemoryMap( - localFragmentInstanceId.queryId, fullFragmentInstanceId, localPlanNodeId); + localFragmentInstanceId.queryId, fullFragmentInstanceId, memoryReservationPlanNodeId); } @Override @@ -271,7 +275,7 @@ public synchronized ByteBuffer getSerializedTsBlock() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, retainedSize); } @@ -305,7 +309,10 @@ private synchronized void trySubmitGetDataBlocksTask() { long reservedBytes = 0L; Pair, Boolean> pair = null; long blockedSize = 0L; - while (sequenceIdToDataBlockSize.containsKey(endSequenceId)) { + // Keep one in-flight block per source handle when it participates in a shared Collect quota. + // Standalone exchanges retain their existing batching behavior. + while (sequenceIdToDataBlockSize.containsKey(endSequenceId) + && (!usesSharedMemoryReservationGroup || endSequenceId == startSequenceId)) { Long bytesToReserve = sequenceIdToDataBlockSize.get(endSequenceId); if (bytesToReserve == null) { throw new IllegalStateException(DataNodeQueryMessages.DATA_BLOCK_SIZE_IS_NULL); @@ -316,7 +323,7 @@ private synchronized void trySubmitGetDataBlocksTask() { .reserveWithPriority( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bytesToReserve, maxBytesCanReserve, isHighestPriority); @@ -446,7 +453,7 @@ public synchronized void abort() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } @@ -474,7 +481,7 @@ public synchronized void abort(Throwable t) { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } @@ -502,7 +509,7 @@ public synchronized void close() { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, bufferRetainedSizeInBytes); bufferRetainedSizeInBytes = 0; } @@ -551,6 +558,23 @@ public void setMaxBytesCanReserve(long maxBytesCanReserve) { this.maxBytesCanReserve = Math.min(this.maxBytesCanReserve, maxBytesCanReserve); } + @Override + public synchronized void setMemoryReservationGroupId(String reservationPlanNodeId) { + this.memoryReservationPlanNodeId = Validate.notNull(reservationPlanNodeId); + this.usesSharedMemoryReservationGroup = !localPlanNodeId.equals(reservationPlanNodeId); + if (usesSharedMemoryReservationGroup) { + // The planner may first apply the legacy per-exchange limit before it knows this handle's + // Collect group. Restore the fragment-instance limit so the shared reservation key, rather + // than an individual handle, is the group-wide limit. + this.maxBytesCanReserve = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + } + localMemoryManager + .getQueryPool() + .registerPlanNodeIdToQueryMemoryMap( + localFragmentInstanceId.queryId, fullFragmentInstanceId, memoryReservationPlanNodeId); + } + @Override public boolean isAborted() { return aborted; @@ -677,6 +701,7 @@ public void run() { indexOfUpstreamSinkHandle); } if (tsBlockNum == 0) { + fail(new IllegalStateException()); return; } throw new TException( @@ -787,7 +812,7 @@ private void fail(Throwable t) { .free( localFragmentInstanceId.getQueryId(), fullFragmentInstanceId, - localPlanNodeId, + memoryReservationPlanNodeId, reservedBytes); } sourceHandleListener.onFailure(SourceHandle.this, t); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/CollectQuotaGroup.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/CollectQuotaGroup.java new file mode 100644 index 0000000000000..69f02fc486dea --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/CollectQuotaGroup.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.db.queryengine.plan.planner; + +import org.apache.iotdb.db.queryengine.execution.exchange.source.ISourceHandle; + +import java.util.ArrayList; +import java.util.List; + +/** The shared memory reservation scope of the exchange inputs of one Collect node. */ +public class CollectQuotaGroup { + private final String reservationPlanNodeId; + private final List sourceHandles = new ArrayList<>(); + + public CollectQuotaGroup(String reservationPlanNodeId) { + this.reservationPlanNodeId = reservationPlanNodeId; + } + + public String getReservationPlanNodeId() { + return reservationPlanNodeId; + } + + public void addSourceHandle(ISourceHandle sourceHandle) { + sourceHandles.add(sourceHandle); + } + + public List getSourceHandles() { + return sourceHandles; + } +} 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 cb2f6f02acb57..157555b442497 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 @@ -47,6 +47,7 @@ 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.sql.ast.Expression; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall; @@ -229,6 +230,16 @@ public DataNodeTableOperatorGenerator(Metadata metadata) { super(metadata); } + @Override + public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context) { + CollectQuotaGroup quotaGroup = context.enterCollectQuotaGroup(node.getPlanNodeId().getId()); + try { + return super.visitCollect(node, context); + } finally { + context.exitCollectQuotaGroup(quotaGroup); + } + } + /** * Producer side: register a shared {@link TopKRuntimeFilter} before visiting TopK children so * nested Scan nodes can resolve the same instance during operator generation. @@ -385,6 +396,7 @@ public Operator visitTableExchange(ExchangeNode node, LocalExecutionPlanContext if (!isSameNode) { context.addExchangeSumNum(1); } + context.registerExchangeSourceHandle(sourceHandle); sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); ExchangeOperator exchangeOperator = new ExchangeOperator(operatorContext, sourceHandle, node.getUpstreamPlanNodeId()); 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 7a5f52749aa4a..198942f06e31c 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,6 +32,7 @@ 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.ISourceHandle; 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; @@ -53,8 +54,12 @@ import javax.annotation.Nullable; import java.time.ZoneId; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Collections; +import java.util.Deque; import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; @@ -83,6 +88,8 @@ public class LocalExecutionPlanContext implements ITableOperatorGeneratorContext private List pipelineDriverFactories; private List exchangeOperatorList = new ArrayList<>(); private int exchangeSumNum = 0; + private List collectQuotaGroups = new ArrayList<>(); + private Deque activeCollectQuotaGroups = new ArrayDeque<>(); private List cachedDataTypes; @@ -129,6 +136,8 @@ public LocalExecutionPlanContext(LocalExecutionPlanContext parentContext) { this.degreeOfParallelism = parentContext.degreeOfParallelism; this.exchangeSumNum = parentContext.exchangeSumNum; this.exchangeOperatorList = parentContext.exchangeOperatorList; + this.collectQuotaGroups = parentContext.collectQuotaGroups; + this.activeCollectQuotaGroups = parentContext.activeCollectQuotaGroups; this.cachedDataTypes = parentContext.cachedDataTypes; this.driverContext = parentContext.getDriverContext().createSubDriverContext(getNextPipelineId()); @@ -261,6 +270,26 @@ public void addExchangeOperator(ExchangeOperator exchangeOperator) { this.exchangeOperatorList.add(exchangeOperator); } + public CollectQuotaGroup enterCollectQuotaGroup(String planNodeId) { + CollectQuotaGroup group = new CollectQuotaGroup("collect-fan-in-" + planNodeId); + collectQuotaGroups.add(group); + activeCollectQuotaGroups.push(group); + return group; + } + + public void exitCollectQuotaGroup(CollectQuotaGroup group) { + if (activeCollectQuotaGroups.peek() != group) { + throw new IllegalStateException(); + } + activeCollectQuotaGroups.pop(); + } + + public void registerExchangeSourceHandle(ISourceHandle sourceHandle) { + if (!activeCollectQuotaGroups.isEmpty()) { + activeCollectQuotaGroups.peek().addSourceHandle(sourceHandle); + } + } + public void setMaxBytesOneHandleCanReserve() { long maxBytesOneHandleCanReserve = getMaxBytesOneHandleCanReserve(); LOGGER.debug( @@ -268,9 +297,23 @@ public void setMaxBytesOneHandleCanReserve() { .MAXBYTESONEHANDLECANRESERVE_FOR_EXCHANGEOPERATOR_IS_ARG_EXCHANGESUMNUM_IS_ARG, maxBytesOneHandleCanReserve, exchangeSumNum); + java.util.Set groupedSourceHandles = + Collections.newSetFromMap(new IdentityHashMap<>()); + long maxBytesPerFI = + IoTDBDescriptor.getInstance().getMemoryConfig().getMaxBytesPerFragmentInstance(); + for (CollectQuotaGroup group : collectQuotaGroups) { + for (ISourceHandle sourceHandle : group.getSourceHandles()) { + groupedSourceHandles.add(sourceHandle); + sourceHandle.setMemoryReservationGroupId(group.getReservationPlanNodeId()); + sourceHandle.setMaxBytesCanReserve(maxBytesPerFI); + } + } exchangeOperatorList.forEach( - exchangeOperator -> - exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve)); + exchangeOperator -> { + if (!groupedSourceHandles.contains(exchangeOperator.getSourceHandle())) { + exchangeOperator.getSourceHandle().setMaxBytesCanReserve(maxBytesOneHandleCanReserve); + } + }); } public Set getAllSensors(IDeviceID deviceId, String sensorId) { 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 b6a57db5445d9..fc57a1764ee78 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 @@ -1181,7 +1181,13 @@ public Operator visitCollect(CollectNode node, LocalExecutionPlanContext context context.getNextOperatorId(), node.getPlanNodeId(), CollectOperator.class.getSimpleName()); - return new CollectOperator(operatorContext, dealWithConsumeChildrenOneByOneNode(node, context)); + CollectQuotaGroup quotaGroup = context.enterCollectQuotaGroup(node.getPlanNodeId().getId()); + try { + return new CollectOperator( + operatorContext, dealWithConsumeChildrenOneByOneNode(node, context)); + } finally { + context.exitCollectQuotaGroup(quotaGroup); + } } @Override @@ -2578,6 +2584,7 @@ public Operator visitExchange(ExchangeNode node, LocalExecutionPlanContext conte context.addExchangeSumNum(1); } sourceHandle.setMaxBytesCanReserve(context.getMaxBytesOneHandleCanReserve()); + context.registerExchangeSourceHandle(sourceHandle); 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/memory/NotThreadSafeMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java index 39e987a4b2f8d..15d64cb97c9c6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/NotThreadSafeMemoryReservationManager.java @@ -143,6 +143,14 @@ public void releaseMemoryImmediately(final long size) { } } + @Override + public void releaseMemoryImmediately() { + if (bytesToBeReleased > 0) { + releaseBytesImmediately(bytesToBeReleased); + bytesToBeReleased = 0; + } + } + private void releaseBytesImmediately(final long size) { long poolBytes = deductReleaseAccounting(size); if (poolBytes > 0) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java index 71676e5b77fe9..e42c9bb87d713 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/memory/ThreadSafeMemoryReservationManager.java @@ -56,6 +56,11 @@ public synchronized void releaseMemoryImmediately(long size) { super.releaseMemoryImmediately(size); } + @Override + public synchronized void releaseMemoryImmediately() { + super.releaseMemoryImmediately(); + } + @Override public synchronized void releaseAllReservedMemory() { super.releaseAllReservedMemory(); 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 ec095730f5417..0a1592b96b63d 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 @@ -2956,6 +2956,12 @@ private static OrderingScheme constructOrderingSchema(List symbols) { return new OrderingScheme(symbols, orderings); } + private static boolean isStreamingAggregationEnabled(AggregationNode node) { + // Keep an explicit ordering boundary for plans whose logical grouping property is streamable. + // The local operator generator still chooses blocking HashAggregation for these inputs. + return node.isStreamable(); + } + private PlanNode mergeChildrenViaCollectOrMergeSort( final OrderingScheme childOrdering, final List childrenNodes) { checkArgument( 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 a92eaa5d0e283..8d7ef53dd1a24 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 @@ -166,6 +166,28 @@ public void concurrencyTest() { } } + @Test + public void groupedQueueAllowsOnlyOneInFlightBlock() { + LocalMemoryManager mockLocalMemoryManager = Mockito.mock(LocalMemoryManager.class); + MemoryPool mockMemoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(mockLocalMemoryManager.getQueryPool()).thenReturn(mockMemoryPool); + SharedTsBlockQueue queue = + new SharedTsBlockQueue( + new TFragmentInstanceId("q0", 0, "0"), + "exchange_0", + mockLocalMemoryManager, + newDirectExecutorService()); + queue.setMemoryReservationGroupId("collect-fan-in-0"); + queue.setMaxBytesCanReserve(Long.MAX_VALUE); + + ListenableFuture producerMayContinue = queue.add(Utils.createMockTsBlock(1024L)); + Assert.assertFalse(producerMayContinue.isDone()); + Assert.assertTrue(queue.isBlocked().isDone()); + + queue.remove(); + Assert.assertTrue(producerMayContinue.isDone()); + } + private static class SendTask implements Runnable { private final SharedTsBlockQueue queue; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/CollectOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/CollectOperatorTest.java new file mode 100644 index 0000000000000..b19dd4a8d46c0 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/CollectOperatorTest.java @@ -0,0 +1,146 @@ +/* + * 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.operator; + +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 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.apache.tsfile.utils.RamUsageEstimator; +import org.junit.Test; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static com.google.common.util.concurrent.Futures.immediateVoidFuture; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +public class CollectOperatorTest { + + @Test + public void shouldConsumeReadyChildBeforeEarlierBlockedChild() throws Exception { + SettableFuture firstChildBlocked = SettableFuture.create(); + TestOperator firstChild = new TestOperator(block(1), firstChildBlocked); + TestOperator secondChild = new TestOperator(block(2), immediateVoidFuture()); + + CollectOperator collectOperator = + new CollectOperator(null, Arrays.asList(firstChild, secondChild)); + + assertTrue(collectOperator.isBlocked().isDone()); + assertSame(secondChild.blocks.get(0), collectOperator.next()); + assertTrue(collectOperator.hasNext()); + } + + @Test + public void shouldRotateReadyChildren() throws Exception { + TestOperator firstChild = + new TestOperator(Arrays.asList(block(1), block(3)), immediateVoidFuture()); + TestOperator secondChild = + new TestOperator(Arrays.asList(block(2), block(4)), immediateVoidFuture()); + + CollectOperator collectOperator = + new CollectOperator(null, Arrays.asList(firstChild, secondChild)); + + assertSame(firstChild.blocks.get(0), collectOperator.next()); + assertSame(secondChild.blocks.get(0), collectOperator.next()); + assertSame(firstChild.blocks.get(1), collectOperator.next()); + assertSame(secondChild.blocks.get(1), collectOperator.next()); + } + + private static TsBlock block(int value) { + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT32)); + builder.getTimeColumnBuilder().writeLong(value); + builder.getColumnBuilder(0).writeInt(value); + builder.declarePosition(); + return builder.build(); + } + + private static final class TestOperator implements Operator { + private final List blocks; + private final ListenableFuture blocked; + private final CommonOperatorContext operatorContext = Mockito.mock(CommonOperatorContext.class); + private int index; + + private TestOperator(TsBlock block, ListenableFuture blocked) { + this(Collections.singletonList(block), blocked); + } + + private TestOperator(List blocks, ListenableFuture blocked) { + this.blocks = blocks; + this.blocked = blocked; + } + + @Override + public CommonOperatorContext getOperatorContext() { + return operatorContext; + } + + @Override + public TsBlock next() { + return blocks.get(index++); + } + + @Override + public boolean hasNext() { + return index < blocks.size(); + } + + @Override + public ListenableFuture isBlocked() { + return blocked; + } + + @Override + public boolean isFinished() { + return !hasNext(); + } + + @Override + public void close() {} + + @Override + public long calculateMaxPeekMemory() { + return 0; + } + + @Override + public long calculateMaxReturnSize() { + return 0; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return 0; + } + + @Override + public long ramBytesUsed() { + return RamUsageEstimator.shallowSizeOfInstance(TestOperator.class); + } + } +}