Implement shared fan-in quota and ready-first Collect

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<Operator> 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<Operator> 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<Void> 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

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<Integer> 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<Integer> 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 {

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);
   }

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();
   }
 

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.